Merge commit '00ad02d04728a8e32a82a4e72747fefcf8a82d4e' into refactor/T-765-session-feature-boundary
# Conflicts: # crates/worker/src/feature/builtin.rs
This commit is contained in:
@@ -475,6 +475,7 @@ struct WorkdirCommandSession {
|
||||
attachment_alias: String,
|
||||
source: WorkdirSessionHandle,
|
||||
provider_handle: CommandHandle,
|
||||
terminal_at: Option<std::time::Instant>,
|
||||
}
|
||||
|
||||
enum RegisteredWorkdirSession {
|
||||
@@ -515,6 +516,7 @@ impl WorkdirSessionRegistry {
|
||||
source: WorkdirSessionHandle,
|
||||
provider_handle: CommandHandle,
|
||||
) -> CommandHandle {
|
||||
self.prune_terminal_commands(&worker);
|
||||
let external_handle = loop {
|
||||
let candidate = CommandHandle(Uuid::now_v7().to_string());
|
||||
if !self
|
||||
@@ -530,6 +532,7 @@ impl WorkdirSessionRegistry {
|
||||
attachment_alias,
|
||||
source,
|
||||
provider_handle,
|
||||
terminal_at: None,
|
||||
},
|
||||
);
|
||||
external_handle
|
||||
@@ -547,6 +550,45 @@ impl WorkdirSessionRegistry {
|
||||
.cloned()
|
||||
}
|
||||
|
||||
fn active_command_count(&self, worker: &RuntimeWorkerRef) -> usize {
|
||||
self.commands
|
||||
.iter()
|
||||
.filter(|((owner, _), command)| owner == worker && command.terminal_at.is_none())
|
||||
.count()
|
||||
}
|
||||
|
||||
fn mark_terminal(&mut self, worker: &RuntimeWorkerRef, external_handle: &CommandHandle) {
|
||||
if let Some(command) = self
|
||||
.commands
|
||||
.get_mut(&(worker.clone(), external_handle.clone()))
|
||||
{
|
||||
command
|
||||
.terminal_at
|
||||
.get_or_insert_with(std::time::Instant::now);
|
||||
}
|
||||
self.prune_terminal_commands(worker);
|
||||
}
|
||||
|
||||
fn prune_terminal_commands(&mut self, worker: &RuntimeWorkerRef) {
|
||||
let mut terminal = self
|
||||
.commands
|
||||
.iter()
|
||||
.filter_map(|((owner, handle), command)| {
|
||||
(owner == worker)
|
||||
.then_some(command.terminal_at)
|
||||
.flatten()
|
||||
.map(|at| (handle.clone(), at))
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
terminal.sort_by_key(|(_, at)| *at);
|
||||
let excess = terminal
|
||||
.len()
|
||||
.saturating_sub(MAX_RETAINED_TERMINAL_COMMANDS);
|
||||
for (handle, _) in terminal.into_iter().take(excess) {
|
||||
self.commands.remove(&(worker.clone(), handle));
|
||||
}
|
||||
}
|
||||
|
||||
fn take_worker_commands_for_alias(
|
||||
&mut self,
|
||||
worker: &RuntimeWorkerRef,
|
||||
@@ -895,7 +937,27 @@ struct ExternalProviderWorkdirSession {
|
||||
audit: Option<ExternalProviderAuditContext>,
|
||||
closed: std::sync::atomic::AtomicBool,
|
||||
close_lock: tokio::sync::Mutex<()>,
|
||||
active_commands: tokio::sync::Mutex<HashMap<String, workdir::CommandHandle>>,
|
||||
active_commands: tokio::sync::Mutex<HashMap<String, Arc<ExternalProviderCommandState>>>,
|
||||
}
|
||||
|
||||
const MAX_RETAINED_TERMINAL_COMMANDS: usize = 128;
|
||||
const MAX_ACTIVE_SERVER_COMMANDS: usize = 128;
|
||||
|
||||
struct ExternalProviderCommandState {
|
||||
handle: workdir::CommandHandle,
|
||||
// Serialize output ACKs, not cancellation: cancel must unblock a waiting reader.
|
||||
cursor: tokio::sync::Mutex<u64>,
|
||||
terminal_at: Mutex<Option<std::time::Instant>>,
|
||||
}
|
||||
|
||||
impl ExternalProviderCommandState {
|
||||
fn new(handle: workdir::CommandHandle) -> Self {
|
||||
Self {
|
||||
handle,
|
||||
cursor: tokio::sync::Mutex::new(0),
|
||||
terminal_at: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ExternalProviderWorkdirSession {
|
||||
@@ -986,17 +1048,24 @@ impl ExternalProviderWorkdirSession {
|
||||
"External Workdir grant has expired".to_string(),
|
||||
));
|
||||
}
|
||||
let _admission = self
|
||||
.connection
|
||||
.admission
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|_| {
|
||||
workdir::WorkdirError::Unavailable(
|
||||
"External Workdir provider admission is closed".to_string(),
|
||||
)
|
||||
})?;
|
||||
// Cancellation must reach the provider even when waiting output readers
|
||||
// occupy every data-operation slot. The socket separately bounds control requests.
|
||||
let _admission = if matches!(&operation, WorkdirSessionOperation::CommandCancel(_)) {
|
||||
None
|
||||
} else {
|
||||
Some(
|
||||
self.connection
|
||||
.admission
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|_| {
|
||||
workdir::WorkdirError::Unavailable(
|
||||
"External Workdir provider admission is closed".to_string(),
|
||||
)
|
||||
})?,
|
||||
)
|
||||
};
|
||||
if self
|
||||
.connection
|
||||
.expires_at
|
||||
@@ -1073,10 +1142,10 @@ impl ExternalProviderWorkdirSession {
|
||||
}
|
||||
}
|
||||
|
||||
async fn provider_command_handle(
|
||||
async fn provider_command(
|
||||
&self,
|
||||
external_handle: &workdir::CommandHandle,
|
||||
) -> std::result::Result<workdir::CommandHandle, workdir::WorkdirError> {
|
||||
) -> std::result::Result<Arc<ExternalProviderCommandState>, workdir::WorkdirError> {
|
||||
if self.closed.load(Ordering::Acquire) {
|
||||
return Err(workdir::WorkdirError::SessionClosed);
|
||||
}
|
||||
@@ -1088,51 +1157,111 @@ impl ExternalProviderWorkdirSession {
|
||||
.ok_or_else(|| workdir::WorkdirError::UnknownCommand(external_handle.0.clone()))
|
||||
}
|
||||
|
||||
async fn provider_command_handle(
|
||||
&self,
|
||||
external_handle: &workdir::CommandHandle,
|
||||
) -> std::result::Result<workdir::CommandHandle, workdir::WorkdirError> {
|
||||
Ok(self.provider_command(external_handle).await?.handle.clone())
|
||||
}
|
||||
|
||||
async fn remove_active_command_if_matches(
|
||||
&self,
|
||||
external_handle: &workdir::CommandHandle,
|
||||
provider_handle: &workdir::CommandHandle,
|
||||
command: &Arc<ExternalProviderCommandState>,
|
||||
) {
|
||||
let mut commands = self.active_commands.lock().await;
|
||||
if commands.get(&external_handle.0) == Some(provider_handle) {
|
||||
if commands
|
||||
.get(&external_handle.0)
|
||||
.is_some_and(|current| Arc::ptr_eq(current, command))
|
||||
{
|
||||
commands.remove(&external_handle.0);
|
||||
}
|
||||
}
|
||||
|
||||
async fn prune_terminal_commands(&self) {
|
||||
let mut commands = self.active_commands.lock().await;
|
||||
let mut terminal = commands
|
||||
.iter()
|
||||
.filter_map(|(key, command)| {
|
||||
command
|
||||
.terminal_at
|
||||
.lock()
|
||||
.unwrap()
|
||||
.map(|at| (key.clone(), at))
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
terminal.sort_by_key(|(_, at)| *at);
|
||||
let excess = terminal
|
||||
.len()
|
||||
.saturating_sub(MAX_RETAINED_TERMINAL_COMMANDS);
|
||||
for (key, _) in terminal.into_iter().take(excess) {
|
||||
commands.remove(&key);
|
||||
}
|
||||
}
|
||||
|
||||
async fn terminate_provider_command(
|
||||
&self,
|
||||
provider_handle: workdir::CommandHandle,
|
||||
command: Arc<ExternalProviderCommandState>,
|
||||
reused_handle: bool,
|
||||
) -> std::result::Result<(), workdir::WorkdirError> {
|
||||
// Cancellation is best effort because the command may already be terminal.
|
||||
// Terminal output is the proof that the provider reaped the process.
|
||||
if !reused_handle && command.terminal_at.lock().unwrap().is_some() {
|
||||
return Ok(());
|
||||
}
|
||||
// Do not acquire the output lock until cancellation has unblocked any reader.
|
||||
// The reader records its last known cursor before cleanup acknowledges bytes.
|
||||
let _ = self
|
||||
.operate(WorkdirSessionOperation::CommandCancel(
|
||||
provider_handle.clone(),
|
||||
command.handle.clone(),
|
||||
))
|
||||
.await;
|
||||
match self
|
||||
.operate(WorkdirSessionOperation::CommandOutput(
|
||||
workdir::CommandOutputRequest {
|
||||
handle: provider_handle,
|
||||
cursor: 0,
|
||||
limit: workdir::external::MAX_EXTERNAL_COMMAND_OUTPUT_BYTES,
|
||||
wait: true,
|
||||
},
|
||||
))
|
||||
.await?
|
||||
{
|
||||
WorkdirSessionOperationResult::CommandOutput(output)
|
||||
if output.status != workdir::CommandStatus::Running =>
|
||||
{
|
||||
Ok(())
|
||||
}
|
||||
WorkdirSessionOperationResult::CommandOutput(_) => {
|
||||
Err(workdir::WorkdirError::Unavailable(
|
||||
"External Workdir provider did not confirm command termination".to_string(),
|
||||
))
|
||||
}
|
||||
_ => Err(Self::mismatch("command_output")),
|
||||
let mut cursor =
|
||||
tokio::time::timeout(std::time::Duration::from_secs(30), command.cursor.lock())
|
||||
.await
|
||||
.map_err(|_| {
|
||||
workdir::WorkdirError::Unavailable(
|
||||
"External Workdir reader did not stop after cancellation".to_string(),
|
||||
)
|
||||
})?;
|
||||
if !reused_handle && command.terminal_at.lock().unwrap().is_some() {
|
||||
return Ok(());
|
||||
}
|
||||
let drain = async {
|
||||
loop {
|
||||
match self
|
||||
.operate(WorkdirSessionOperation::CommandOutput(
|
||||
workdir::CommandOutputRequest {
|
||||
handle: command.handle.clone(),
|
||||
cursor: *cursor,
|
||||
limit: workdir::MAX_COMMAND_READ_BYTES,
|
||||
wait: true,
|
||||
},
|
||||
))
|
||||
.await?
|
||||
{
|
||||
WorkdirSessionOperationResult::CommandOutput(output) => {
|
||||
*cursor = output.next_cursor;
|
||||
if output.terminal.is_some() {
|
||||
*command.terminal_at.lock().unwrap() = Some(std::time::Instant::now());
|
||||
return Ok(());
|
||||
}
|
||||
if output.chunks.is_empty() {
|
||||
return Err(workdir::WorkdirError::Unavailable(
|
||||
"External Workdir provider did not confirm command termination"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
_ => return Err(Self::mismatch("command_output")),
|
||||
}
|
||||
}
|
||||
};
|
||||
tokio::time::timeout(std::time::Duration::from_secs(30), drain)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
workdir::WorkdirError::Unavailable(
|
||||
"External Workdir command drain timed out".to_string(),
|
||||
)
|
||||
})?
|
||||
}
|
||||
|
||||
fn mismatch(expected: &str) -> workdir::WorkdirError {
|
||||
@@ -1410,17 +1539,27 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession {
|
||||
}
|
||||
async fn start_command(
|
||||
&self,
|
||||
mut request: workdir::CommandRequest,
|
||||
request: workdir::CommandRequest,
|
||||
) -> std::result::Result<workdir::CommandHandle, workdir::WorkdirError> {
|
||||
self.ensure_capability(workdir::WorkdirSessionCapability::Command)?;
|
||||
let _close_guard = self.close_lock.lock().await;
|
||||
if self.closed.load(Ordering::Acquire) {
|
||||
return Err(workdir::WorkdirError::SessionClosed);
|
||||
}
|
||||
// Bash supplies a Worker-host spill directory. It is neither meaningful
|
||||
// nor safe to send to a client-hosted provider, so External commands use
|
||||
// bounded inline output and report truncation without a provider path.
|
||||
request.spill_dir = None;
|
||||
self.prune_terminal_commands().await;
|
||||
if self
|
||||
.active_commands
|
||||
.lock()
|
||||
.await
|
||||
.values()
|
||||
.filter(|command| command.terminal_at.lock().unwrap().is_none())
|
||||
.count()
|
||||
>= MAX_ACTIVE_SERVER_COMMANDS
|
||||
{
|
||||
return Err(workdir::WorkdirError::Unavailable(
|
||||
"External Workdir active command limit reached".to_string(),
|
||||
));
|
||||
}
|
||||
match self
|
||||
.operate(WorkdirSessionOperation::CommandStart(request))
|
||||
.await?
|
||||
@@ -1430,7 +1569,7 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession {
|
||||
let commands = self.active_commands.lock().await;
|
||||
commands
|
||||
.iter()
|
||||
.filter(|(_, active)| *active == &provider_handle)
|
||||
.filter(|(_, active)| active.handle == provider_handle)
|
||||
.map(|(external, _)| external.clone())
|
||||
.collect::<Vec<_>>()
|
||||
};
|
||||
@@ -1439,20 +1578,29 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession {
|
||||
// generation. Remove every ambiguous route before attempting
|
||||
// cleanup so no caller can address the wrong process.
|
||||
let mut commands = self.active_commands.lock().await;
|
||||
let mut quarantined = None;
|
||||
for external_handle in duplicate_handles {
|
||||
commands.remove(&external_handle);
|
||||
if let Some(command) = commands.remove(&external_handle) {
|
||||
quarantined.get_or_insert(command);
|
||||
}
|
||||
}
|
||||
let quarantine_handle =
|
||||
workdir::CommandHandle(format!("quarantined-command-{}", Uuid::now_v7()));
|
||||
commands.insert(quarantine_handle.0.clone(), provider_handle.clone());
|
||||
let command = quarantined.expect("duplicate command route disappeared");
|
||||
*command.terminal_at.lock().unwrap() = None;
|
||||
commands.insert(quarantine_handle.0.clone(), command.clone());
|
||||
drop(commands);
|
||||
if self
|
||||
.terminate_provider_command(provider_handle.clone())
|
||||
.terminate_provider_command(command.clone(), true)
|
||||
.await
|
||||
.is_ok()
|
||||
{
|
||||
self.remove_active_command_if_matches(&quarantine_handle, &provider_handle)
|
||||
self.remove_active_command_if_matches(&quarantine_handle, &command)
|
||||
.await;
|
||||
} else {
|
||||
// An earlier read of the reused handle cannot establish that
|
||||
// the newly started, ambiguous process has been cleaned up.
|
||||
*command.terminal_at.lock().unwrap() = None;
|
||||
}
|
||||
return Err(workdir::WorkdirError::Unavailable(
|
||||
"External Workdir provider reused an active command handle".to_string(),
|
||||
@@ -1460,10 +1608,10 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession {
|
||||
}
|
||||
let external_handle =
|
||||
workdir::CommandHandle(format!("external-command-{}", Uuid::now_v7()));
|
||||
self.active_commands
|
||||
.lock()
|
||||
.await
|
||||
.insert(external_handle.0.clone(), provider_handle);
|
||||
self.active_commands.lock().await.insert(
|
||||
external_handle.0.clone(),
|
||||
Arc::new(ExternalProviderCommandState::new(provider_handle)),
|
||||
);
|
||||
Ok(external_handle)
|
||||
}
|
||||
_ => Err(Self::mismatch("command_start")),
|
||||
@@ -1489,21 +1637,29 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession {
|
||||
) -> std::result::Result<workdir::CommandOutput, workdir::WorkdirError> {
|
||||
self.ensure_capability(workdir::WorkdirSessionCapability::Command)?;
|
||||
let external_handle = request.handle.clone();
|
||||
let provider_handle = self.provider_command_handle(&external_handle).await?;
|
||||
request.handle = provider_handle.clone();
|
||||
let command = self.provider_command(&external_handle).await?;
|
||||
let mut cursor = command.cursor.lock().await;
|
||||
request.handle = command.handle.clone();
|
||||
// If the response is lost, this request cursor still replays the provider page.
|
||||
match self
|
||||
.operate(WorkdirSessionOperation::CommandOutput(request))
|
||||
.await?
|
||||
.await
|
||||
{
|
||||
WorkdirSessionOperationResult::CommandOutput(output) => {
|
||||
if output.status != workdir::CommandStatus::Running && output.next_cursor.is_none()
|
||||
{
|
||||
self.remove_active_command_if_matches(&external_handle, &provider_handle)
|
||||
.await;
|
||||
Ok(WorkdirSessionOperationResult::CommandOutput(output)) => {
|
||||
*cursor = output.next_cursor;
|
||||
if output.terminal.is_some() {
|
||||
command
|
||||
.terminal_at
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get_or_insert_with(std::time::Instant::now);
|
||||
}
|
||||
drop(cursor);
|
||||
self.prune_terminal_commands().await;
|
||||
Ok(output)
|
||||
}
|
||||
_ => Err(Self::mismatch("command_output")),
|
||||
Ok(_) => Err(Self::mismatch("command_output")),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
async fn cancel_command(
|
||||
@@ -1540,21 +1696,25 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession {
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let mut cleanup_errors = Vec::new();
|
||||
for (external_handle, provider_handle) in commands {
|
||||
for (external_handle, command) in commands {
|
||||
// Detach/session replacement is terminal authority for commands
|
||||
// started through this exact provider generation. Do not report a
|
||||
// successful close until the provider confirms terminal output.
|
||||
match self
|
||||
.terminate_provider_command(provider_handle.clone())
|
||||
.terminate_provider_command(command.clone(), false)
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
self.remove_active_command_if_matches(&external_handle, &provider_handle)
|
||||
self.remove_active_command_if_matches(&external_handle, &command)
|
||||
.await;
|
||||
}
|
||||
Err(error) => {
|
||||
let still_active = self.active_commands.lock().await.get(&external_handle.0)
|
||||
== Some(&provider_handle);
|
||||
let still_active = self
|
||||
.active_commands
|
||||
.lock()
|
||||
.await
|
||||
.get(&external_handle.0)
|
||||
.is_some_and(|current| Arc::ptr_eq(current, &command));
|
||||
if still_active {
|
||||
cleanup_errors.push(error.to_string());
|
||||
}
|
||||
@@ -12534,6 +12694,14 @@ impl server_api::ServerApi for ServerApiContractService {
|
||||
let api = self.workspace_api().map_err(|_| drive::failure(404))?;
|
||||
drive::root(api, &context, &workspace_id).await
|
||||
}
|
||||
async fn drive_capabilities(
|
||||
&self,
|
||||
context: server_api::ServerRequestContext,
|
||||
workspace_id: String,
|
||||
) -> std::result::Result<server_api::DriveCapabilities, server_api::DriveApiError> {
|
||||
let api = self.workspace_api().map_err(|_| drive::failure(404))?;
|
||||
drive::capabilities(api, &context, &workspace_id).await
|
||||
}
|
||||
async fn drive_metadata(
|
||||
&self,
|
||||
context: server_api::ServerRequestContext,
|
||||
@@ -20811,6 +20979,9 @@ async fn scoped_execute_current_worker_workdir_operation(
|
||||
&api, &worker, &target_workdir, Some(&_session_guard),
|
||||
).await?;
|
||||
context.workdir_id_hash = Some(crate::worker_source::diagnostic_id_hash(&link.workdir_id));
|
||||
if api.workdir_sessions.lock().unwrap().active_command_count(&worker) >= MAX_ACTIVE_SERVER_COMMANDS {
|
||||
return Err(InternalWorkdirOperationError::Provider(workdir_operation_transport_error(workdir::WorkdirError::Unavailable("Worker active command limit reached".to_string()))));
|
||||
}
|
||||
context.stage = "session_open";
|
||||
let source = open_current_worker_workdir_session_locked(&api, &worker, &link, &_session_guard).await?;
|
||||
context.stage = "provider_dispatch";
|
||||
@@ -20852,14 +21023,16 @@ async fn scoped_execute_current_worker_workdir_operation(
|
||||
.map_err(workdir_operation_transport_error)?
|
||||
}
|
||||
WorkdirSessionOperation::CommandOutput(mut output) => {
|
||||
let external_handle = output.handle.clone();
|
||||
let (session, provider_handle) =
|
||||
current_worker_command_session(&api, &worker, &target_workdir, &output.handle, &mut context).await?;
|
||||
current_worker_command_session(&api, &worker, &target_workdir, &external_handle, &mut context).await?;
|
||||
output.handle = provider_handle;
|
||||
session
|
||||
.command_output(output)
|
||||
.await
|
||||
.map(WorkdirSessionOperationResult::CommandOutput)
|
||||
.map_err(workdir_operation_transport_error)?
|
||||
let page = session.command_output(output).await.map_err(workdir_operation_transport_error)?;
|
||||
if page.terminal.is_some() {
|
||||
// Retain the binding for a retry after an unknown HTTP outcome.
|
||||
api.workdir_sessions.lock().unwrap().mark_terminal(&worker, &external_handle);
|
||||
}
|
||||
WorkdirSessionOperationResult::CommandOutput(page)
|
||||
}
|
||||
WorkdirSessionOperation::CommandCancel(external_handle) => {
|
||||
let (session, provider_handle) =
|
||||
@@ -26838,7 +27011,12 @@ async fn serve_external_workdir_provider(
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
if pending.len() >= 16 {
|
||||
let pending_limit = if matches!(&operation, WorkdirSessionOperation::CommandCancel(_)) {
|
||||
16 + MAX_ACTIVE_SERVER_COMMANDS
|
||||
} else {
|
||||
16
|
||||
};
|
||||
if pending.len() >= pending_limit {
|
||||
let _ = response.send(Err(WorkdirTransportError {
|
||||
denial_reason: None,
|
||||
code: workdir::http::WorkdirTransportErrorCode::Unavailable,
|
||||
@@ -37009,6 +37187,7 @@ mod tests {
|
||||
|
||||
mod api_logging_tests;
|
||||
mod auth_logging_tests;
|
||||
mod command_transport_tests;
|
||||
mod drive_api_tests;
|
||||
mod drive_tests;
|
||||
mod drive_web_roundtrip_tests;
|
||||
@@ -37358,9 +37537,8 @@ mod tests {
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "touch forbidden".to_string(),
|
||||
timeout_secs: 1,
|
||||
output_limit: 1024,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: None,
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await;
|
||||
@@ -37509,9 +37687,8 @@ mod tests {
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "printf forbidden".to_string(),
|
||||
timeout_secs: 30,
|
||||
output_limit: 1024,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: None,
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await;
|
||||
@@ -37563,9 +37740,8 @@ mod tests {
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "printf routed".to_string(),
|
||||
timeout_secs: 30,
|
||||
output_limit: 1024,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: Some("/worker/private/spill".into()),
|
||||
tool_call_id: Some("call-1".to_string()),
|
||||
})
|
||||
.await
|
||||
@@ -37581,10 +37757,7 @@ mod tests {
|
||||
};
|
||||
assert!(matches!(
|
||||
operation,
|
||||
WorkdirSessionOperation::CommandStart(workdir::CommandRequest {
|
||||
spill_dir: None,
|
||||
..
|
||||
})
|
||||
WorkdirSessionOperation::CommandStart(workdir::CommandRequest { .. })
|
||||
));
|
||||
let provider_handle = workdir::CommandHandle("command-1".to_string());
|
||||
response
|
||||
@@ -37637,13 +37810,9 @@ mod tests {
|
||||
response
|
||||
.send(Ok(WorkdirSessionOperationResult::CommandOutput(
|
||||
workdir::CommandOutput {
|
||||
status: workdir::CommandStatus::Running,
|
||||
exit_code: None,
|
||||
timed_out: false,
|
||||
content: String::new(),
|
||||
next_cursor: None,
|
||||
truncated: false,
|
||||
output_path: None,
|
||||
chunks: Vec::new(),
|
||||
next_cursor: 0,
|
||||
terminal: None,
|
||||
},
|
||||
)))
|
||||
.unwrap();
|
||||
@@ -37716,13 +37885,22 @@ mod tests {
|
||||
response
|
||||
.send(Ok(WorkdirSessionOperationResult::CommandOutput(
|
||||
workdir::CommandOutput {
|
||||
status: workdir::CommandStatus::Cancelled,
|
||||
exit_code: None,
|
||||
timed_out: false,
|
||||
content: String::new(),
|
||||
next_cursor: None,
|
||||
truncated: false,
|
||||
output_path: None,
|
||||
chunks: Vec::new(),
|
||||
next_cursor: 0,
|
||||
terminal: Some(workdir::CommandTerminal {
|
||||
status: workdir::CommandStatus::Cancelled,
|
||||
exit_code: None,
|
||||
|
||||
stdout_bytes: 0,
|
||||
|
||||
stderr_bytes: 0,
|
||||
|
||||
stdout_delivered_bytes: 0,
|
||||
|
||||
stderr_delivered_bytes: 0,
|
||||
|
||||
truncation_reason: None,
|
||||
}),
|
||||
},
|
||||
)))
|
||||
.unwrap();
|
||||
@@ -37751,9 +37929,8 @@ mod tests {
|
||||
let request = || workdir::CommandRequest {
|
||||
command: "sleep 30".to_string(),
|
||||
timeout_secs: 30,
|
||||
output_limit: 1024,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: None,
|
||||
tool_call_id: None,
|
||||
};
|
||||
let provider_handle = workdir::CommandHandle("command-reused".to_string());
|
||||
@@ -37810,13 +37987,22 @@ mod tests {
|
||||
response
|
||||
.send(Ok(WorkdirSessionOperationResult::CommandOutput(
|
||||
workdir::CommandOutput {
|
||||
status: workdir::CommandStatus::Cancelled,
|
||||
exit_code: None,
|
||||
timed_out: false,
|
||||
content: String::new(),
|
||||
next_cursor: None,
|
||||
truncated: false,
|
||||
output_path: None,
|
||||
chunks: Vec::new(),
|
||||
next_cursor: 0,
|
||||
terminal: Some(workdir::CommandTerminal {
|
||||
status: workdir::CommandStatus::Cancelled,
|
||||
exit_code: None,
|
||||
|
||||
stdout_bytes: 0,
|
||||
|
||||
stderr_bytes: 0,
|
||||
|
||||
stdout_delivered_bytes: 0,
|
||||
|
||||
stderr_delivered_bytes: 0,
|
||||
|
||||
truncation_reason: None,
|
||||
}),
|
||||
},
|
||||
)))
|
||||
.unwrap();
|
||||
@@ -37867,9 +38053,8 @@ mod tests {
|
||||
let request = || workdir::CommandRequest {
|
||||
command: "sleep 30".to_string(),
|
||||
timeout_secs: 30,
|
||||
output_limit: 1024,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: None,
|
||||
tool_call_id: None,
|
||||
};
|
||||
let provider_handle = workdir::CommandHandle("command-1".to_string());
|
||||
@@ -38766,9 +38951,8 @@ mod tests {
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "printf forbidden".to_string(),
|
||||
timeout_secs: 30,
|
||||
output_limit: 1024,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: None,
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await;
|
||||
@@ -38789,9 +38973,8 @@ mod tests {
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "printf attached".to_string(),
|
||||
timeout_secs: 30,
|
||||
output_limit: 1024,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: None,
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await
|
||||
@@ -38894,13 +39077,22 @@ mod tests {
|
||||
response
|
||||
.send(Ok(WorkdirSessionOperationResult::CommandOutput(
|
||||
workdir::CommandOutput {
|
||||
status: workdir::CommandStatus::Cancelled,
|
||||
exit_code: None,
|
||||
timed_out: false,
|
||||
content: String::new(),
|
||||
next_cursor: None,
|
||||
truncated: false,
|
||||
output_path: None,
|
||||
chunks: Vec::new(),
|
||||
next_cursor: 0,
|
||||
terminal: Some(workdir::CommandTerminal {
|
||||
status: workdir::CommandStatus::Cancelled,
|
||||
exit_code: None,
|
||||
|
||||
stdout_bytes: 0,
|
||||
|
||||
stderr_bytes: 0,
|
||||
|
||||
stdout_delivered_bytes: 0,
|
||||
|
||||
stderr_delivered_bytes: 0,
|
||||
|
||||
truncation_reason: None,
|
||||
}),
|
||||
},
|
||||
)))
|
||||
.unwrap();
|
||||
@@ -38985,11 +39177,9 @@ mod tests {
|
||||
};
|
||||
|
||||
let (mut old_provider, _) = connect_async(&url).await.unwrap();
|
||||
let old_registration = serde_json::to_string(®ister(1)).unwrap().replacen(
|
||||
r#""version":2"#,
|
||||
r#""version":1"#,
|
||||
1,
|
||||
);
|
||||
let mut old_registration = serde_json::to_value(register(1)).unwrap();
|
||||
old_registration["version"] = serde_json::json!(1);
|
||||
let old_registration = serde_json::to_string(&old_registration).unwrap();
|
||||
old_provider
|
||||
.send(Message::Text(old_registration.into()))
|
||||
.await
|
||||
@@ -39003,7 +39193,10 @@ mod tests {
|
||||
.reason
|
||||
.contains("unsupported External Workdir protocol version 1")
|
||||
);
|
||||
assert!(frame.reason.contains("requires 2"));
|
||||
assert!(frame.reason.contains(&format!(
|
||||
"requires {}",
|
||||
workdir::external::EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION
|
||||
)));
|
||||
|
||||
let (mut smaller_provider, _) = connect_async(&url).await.unwrap();
|
||||
let mut smaller_registration = serde_json::to_value(register(1)).unwrap();
|
||||
@@ -39244,37 +39437,106 @@ mod tests {
|
||||
fs::read_to_string(provider_root.path().join("created.txt")).unwrap(),
|
||||
"created"
|
||||
);
|
||||
let binary_handle = broker_session
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "head -c 196609 /dev/zero".to_string(),
|
||||
timeout_secs: 30,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let (terminal_request, terminal_page) =
|
||||
tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||||
let mut cursor = 0;
|
||||
let mut bytes = Vec::new();
|
||||
loop {
|
||||
let request = workdir::CommandOutputRequest {
|
||||
handle: binary_handle.clone(),
|
||||
cursor,
|
||||
limit: workdir::MAX_COMMAND_READ_BYTES,
|
||||
wait: true,
|
||||
};
|
||||
let page = broker_session
|
||||
.command_output(request.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
for chunk in &page.chunks {
|
||||
assert_eq!(chunk.stream, workdir::CommandStream::Stdout);
|
||||
assert_eq!(chunk.sequence, bytes.len() as u64);
|
||||
assert_eq!(chunk.offset, bytes.len() as u64);
|
||||
bytes.extend_from_slice(&chunk.bytes);
|
||||
}
|
||||
assert!(
|
||||
page.chunks
|
||||
.iter()
|
||||
.map(|chunk| chunk.bytes.len())
|
||||
.sum::<usize>()
|
||||
<= workdir::MAX_COMMAND_READ_BYTES
|
||||
);
|
||||
cursor = page.next_cursor;
|
||||
if let Some(terminal) = &page.terminal {
|
||||
assert_eq!(terminal.status, workdir::CommandStatus::Completed);
|
||||
assert_eq!(terminal.exit_code, Some(0));
|
||||
assert_eq!(bytes, vec![0; 196609]);
|
||||
break (request, page);
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("incremental provider output did not terminate");
|
||||
assert_eq!(
|
||||
broker_session
|
||||
.command_output(terminal_request.clone())
|
||||
.await
|
||||
.unwrap(),
|
||||
terminal_page
|
||||
);
|
||||
for cursor in [0, terminal_page.next_cursor + 1] {
|
||||
assert!(
|
||||
broker_session
|
||||
.command_output(workdir::CommandOutputRequest {
|
||||
cursor,
|
||||
..terminal_request.clone()
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
let command_handle = broker_session
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "sleep 30 & child=$!; echo $child > child.pid; wait $child".to_string(),
|
||||
command: "sleep 30 & child=$!; echo $child > child.pid; printf ready; wait $child"
|
||||
.to_string(),
|
||||
timeout_secs: 30,
|
||||
output_limit: 1024,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: Some("/worker/private/spill".into()),
|
||||
tool_call_id: Some("call-e2e".to_string()),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
for _ in 0..100 {
|
||||
if provider_root.path().join("child.pid").is_file() {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
assert!(provider_root.path().join("child.pid").is_file());
|
||||
|
||||
let command_output = broker_session
|
||||
.command_output(workdir::CommandOutputRequest {
|
||||
handle: command_handle,
|
||||
cursor: 0,
|
||||
limit: 1024,
|
||||
wait: false,
|
||||
wait: true,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(command_output.status, workdir::CommandStatus::Running);
|
||||
assert_eq!(command_output.output_path, None);
|
||||
assert!(command_output.terminal.is_none());
|
||||
|
||||
assert_eq!(
|
||||
command_output
|
||||
.chunks
|
||||
.iter()
|
||||
.flat_map(|chunk| chunk.bytes.clone())
|
||||
.collect::<Vec<_>>(),
|
||||
b"ready"
|
||||
);
|
||||
assert!(provider_root.path().join("child.pid").is_file());
|
||||
|
||||
let Json(revoked) = scoped_revoke_external_workdir_grant(
|
||||
State(api.clone()),
|
||||
@@ -40477,9 +40739,8 @@ mod tests {
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "printf ready; sleep 30".to_string(),
|
||||
timeout_secs: 60,
|
||||
output_limit: 4096,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: None,
|
||||
tool_call_id: Some("tool-call-command-session".to_string()),
|
||||
})
|
||||
.await
|
||||
@@ -40526,7 +40787,7 @@ mod tests {
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(output.status, workdir::CommandStatus::Running));
|
||||
assert!(output.terminal.is_none());
|
||||
assert!(matches!(
|
||||
command
|
||||
.source
|
||||
@@ -40538,14 +40799,21 @@ mod tests {
|
||||
let waiting_source = command.source.clone();
|
||||
let waiting_handle = command.provider_handle.clone();
|
||||
let waiting = tokio::spawn(async move {
|
||||
waiting_source
|
||||
.command_output(workdir::CommandOutputRequest {
|
||||
handle: waiting_handle,
|
||||
cursor: output.next_cursor.unwrap_or(0),
|
||||
limit: 4096,
|
||||
wait: true,
|
||||
})
|
||||
.await
|
||||
let mut cursor = output.next_cursor;
|
||||
loop {
|
||||
let page = waiting_source
|
||||
.command_output(workdir::CommandOutputRequest {
|
||||
handle: waiting_handle.clone(),
|
||||
cursor,
|
||||
limit: 4096,
|
||||
wait: true,
|
||||
})
|
||||
.await?;
|
||||
cursor = page.next_cursor;
|
||||
if page.terminal.is_some() {
|
||||
return Ok::<_, workdir::WorkdirError>(page);
|
||||
}
|
||||
}
|
||||
});
|
||||
tokio::task::yield_now().await;
|
||||
command
|
||||
@@ -40558,7 +40826,7 @@ mod tests {
|
||||
.expect("cancel should unblock a waiting output read")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(!matches!(terminal.status, workdir::CommandStatus::Running));
|
||||
assert!(terminal.terminal.is_some());
|
||||
|
||||
let registered = registry.take_worker(&worker);
|
||||
assert_eq!(registered.len(), 1);
|
||||
@@ -40591,9 +40859,8 @@ mod tests {
|
||||
let request = || workdir::CommandRequest {
|
||||
command: "sleep 30".to_string(),
|
||||
timeout_secs: 60,
|
||||
output_limit: 4096,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: None,
|
||||
tool_call_id: None,
|
||||
};
|
||||
let left_provider = left.start_command(request()).await.unwrap();
|
||||
@@ -60067,9 +60334,8 @@ mod tests {
|
||||
.start_command(workdir::CommandRequest {
|
||||
command: "sleep 30".to_string(),
|
||||
timeout_secs: 60,
|
||||
output_limit: 4096,
|
||||
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
|
||||
cwd: workdir::WorkdirPath::root(),
|
||||
spill_dir: None,
|
||||
tool_call_id: None,
|
||||
})
|
||||
.await
|
||||
|
||||
@@ -471,11 +471,10 @@ async fn generated_workdir_provider_errors_preserve_response_and_log_details() {
|
||||
));
|
||||
// Seed an already registered command, without launching a real process.
|
||||
let provider_handle = CommandHandle("logging-command".to_string());
|
||||
session
|
||||
.active_commands
|
||||
.lock()
|
||||
.await
|
||||
.insert(provider_handle.0.clone(), provider_handle.clone());
|
||||
session.active_commands.lock().await.insert(
|
||||
provider_handle.0.clone(),
|
||||
Arc::new(ExternalProviderCommandState::new(provider_handle.clone())),
|
||||
);
|
||||
let handle = fixture
|
||||
.api
|
||||
.workdir_sessions
|
||||
@@ -720,7 +719,7 @@ async fn self_workdir_provider_refusals_log_operation_identity_without_request_t
|
||||
Reason::ScopedCapabilityDenied,
|
||||
),
|
||||
(
|
||||
json!({"operation": "command_start", "request": {"command": secret, "timeout_secs": 1, "output_limit": 1024, "cwd": "", "tool_call_id": secret}}),
|
||||
json!({"operation": "command_start", "request": {"command": secret, "timeout_secs": 1, "transport_byte_budget": 1048576, "cwd": "", "tool_call_id": secret}}),
|
||||
"command_start",
|
||||
"denied",
|
||||
Reason::ScopedCapabilityDenied,
|
||||
@@ -942,8 +941,8 @@ async fn self_workdir_post_start_registration_failure_does_not_claim_pre_dispatc
|
||||
serde_json::to_vec(&json!({
|
||||
"target_workdir": "checkout",
|
||||
"operation": {"operation": "command_start", "request": {
|
||||
"command": "body-secret", "timeout_secs": 1, "output_limit": 1024,
|
||||
"cwd": "", "spill_dir": null, "tool_call_id": "body-secret"
|
||||
"command": "body-secret", "timeout_secs": 1, "transport_byte_budget": 1048576,
|
||||
"cwd": "", "tool_call_id": "body-secret"
|
||||
}}
|
||||
}))
|
||||
.unwrap(),
|
||||
@@ -1071,8 +1070,8 @@ async fn saturated_host_before_command_start_returns_availability_without_a_prov
|
||||
serde_json::to_vec(&json!({
|
||||
"target_workdir": "checkout",
|
||||
"operation": {"operation": "command_start", "request": {
|
||||
"command": "not-dispatched", "timeout_secs": 1, "output_limit": 1024,
|
||||
"cwd": "", "spill_dir": null, "tool_call_id": "not-dispatched"
|
||||
"command": "not-dispatched", "timeout_secs": 1, "transport_byte_budget": 1048576,
|
||||
"cwd": "", "tool_call_id": "not-dispatched"
|
||||
}}
|
||||
}))
|
||||
.unwrap(),
|
||||
|
||||
@@ -0,0 +1,223 @@
|
||||
//! Provider-bridge lifecycle tests. The production bridge runs against a controlled
|
||||
//! provider boundary; the websocket/CLI-pump test in the parent module covers real byte transport.
|
||||
use super::*;
|
||||
use workdir::WorkdirSession;
|
||||
|
||||
fn command_fixture() -> (
|
||||
Arc<ExternalProviderWorkdirSession>,
|
||||
tokio::sync::mpsc::Receiver<ExternalProviderCommand>,
|
||||
) {
|
||||
let (sender, receiver) = tokio::sync::mpsc::channel(16);
|
||||
let session = Arc::new(ExternalProviderWorkdirSession::new(
|
||||
Arc::new(ExternalProviderConnection {
|
||||
grant_id: "command-grant".into(),
|
||||
workdir_id: "command-workdir".into(),
|
||||
provider_instance_id: "command-provider".into(),
|
||||
generation: 1,
|
||||
expires_at: None,
|
||||
capabilities: workdir::WorkdirSessionCapabilities::COMMAND_ONLY,
|
||||
admission: Arc::new(tokio::sync::Semaphore::new(1)),
|
||||
shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)),
|
||||
sender,
|
||||
}),
|
||||
None,
|
||||
));
|
||||
(session, receiver)
|
||||
}
|
||||
|
||||
fn page(cursor: u64, bytes: &[u8], terminal: bool) -> workdir::CommandOutput {
|
||||
workdir::CommandOutput {
|
||||
chunks: if bytes.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
vec![workdir::CommandByteChunk {
|
||||
sequence: cursor,
|
||||
stream: workdir::CommandStream::Stdout,
|
||||
offset: cursor,
|
||||
bytes: bytes.to_vec(),
|
||||
}]
|
||||
},
|
||||
next_cursor: cursor + bytes.len() as u64,
|
||||
terminal: terminal.then_some(workdir::CommandTerminal {
|
||||
status: workdir::CommandStatus::Cancelled,
|
||||
exit_code: None,
|
||||
stdout_bytes: cursor + bytes.len() as u64,
|
||||
stderr_bytes: 0,
|
||||
stdout_delivered_bytes: cursor + bytes.len() as u64,
|
||||
stderr_delivered_bytes: 0,
|
||||
truncation_reason: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
async fn receive(
|
||||
receiver: &mut tokio::sync::mpsc::Receiver<ExternalProviderCommand>,
|
||||
) -> (
|
||||
WorkdirSessionOperation,
|
||||
tokio::sync::oneshot::Sender<
|
||||
std::result::Result<WorkdirSessionOperationResult, WorkdirTransportError>,
|
||||
>,
|
||||
) {
|
||||
let command = tokio::time::timeout(std::time::Duration::from_secs(2), receiver.recv())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let ExternalProviderCommand::Operation {
|
||||
operation,
|
||||
response,
|
||||
..
|
||||
} = command
|
||||
else {
|
||||
panic!("expected provider operation")
|
||||
};
|
||||
(operation, response)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn close_cancels_waiting_reader_then_drains_from_latest_cursor() {
|
||||
let (session, mut receiver) = command_fixture();
|
||||
session.active_commands.lock().await.insert(
|
||||
"external".into(),
|
||||
Arc::new(ExternalProviderCommandState::new(CommandHandle(
|
||||
"provider".into(),
|
||||
))),
|
||||
);
|
||||
let reader = {
|
||||
let session = session.clone();
|
||||
tokio::spawn(async move {
|
||||
session
|
||||
.command_output(workdir::CommandOutputRequest {
|
||||
handle: CommandHandle("external".into()),
|
||||
cursor: 0,
|
||||
limit: 3,
|
||||
wait: true,
|
||||
})
|
||||
.await
|
||||
})
|
||||
};
|
||||
let (_, response) = receive(&mut receiver).await;
|
||||
response
|
||||
.send(Ok(WorkdirSessionOperationResult::CommandOutput(page(
|
||||
0, b"one", false,
|
||||
))))
|
||||
.unwrap();
|
||||
assert_eq!(reader.await.unwrap().unwrap().next_cursor, 3);
|
||||
|
||||
let reader = {
|
||||
let session = session.clone();
|
||||
tokio::spawn(async move {
|
||||
session
|
||||
.command_output(workdir::CommandOutputRequest {
|
||||
handle: CommandHandle("external".into()),
|
||||
cursor: 3,
|
||||
limit: 3,
|
||||
wait: true,
|
||||
})
|
||||
.await
|
||||
})
|
||||
};
|
||||
let (_, pending_read) = receive(&mut receiver).await;
|
||||
let closer = {
|
||||
let session = session.clone();
|
||||
tokio::spawn(async move { session.close().await })
|
||||
};
|
||||
let (operation, response) = receive(&mut receiver).await;
|
||||
assert!(
|
||||
matches!(operation, WorkdirSessionOperation::CommandCancel(ref handle) if handle.0 == "provider")
|
||||
);
|
||||
response
|
||||
.send(Ok(WorkdirSessionOperationResult::CommandCancel))
|
||||
.unwrap();
|
||||
pending_read
|
||||
.send(Ok(WorkdirSessionOperationResult::CommandOutput(page(
|
||||
3, b"two", false,
|
||||
))))
|
||||
.unwrap();
|
||||
assert_eq!(reader.await.unwrap().unwrap().next_cursor, 6);
|
||||
|
||||
for (cursor, bytes, terminal) in [(6, &b"end"[..], false), (9, &b""[..], true)] {
|
||||
let (operation, response) = receive(&mut receiver).await;
|
||||
assert!(
|
||||
matches!(operation, WorkdirSessionOperation::CommandOutput(ref request) if request.handle.0 == "provider" && request.cursor == cursor && request.wait)
|
||||
);
|
||||
response
|
||||
.send(Ok(WorkdirSessionOperationResult::CommandOutput(page(
|
||||
cursor, bytes, terminal,
|
||||
))))
|
||||
.unwrap();
|
||||
}
|
||||
tokio::time::timeout(std::time::Duration::from_secs(2), closer)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_output_retains_same_binding_for_current_page_replay() {
|
||||
let (session, mut receiver) = command_fixture();
|
||||
session.active_commands.lock().await.insert(
|
||||
"external".into(),
|
||||
Arc::new(ExternalProviderCommandState::new(CommandHandle(
|
||||
"provider".into(),
|
||||
))),
|
||||
);
|
||||
for _ in 0..2 {
|
||||
let reader = {
|
||||
let session = session.clone();
|
||||
tokio::spawn(async move {
|
||||
session
|
||||
.command_output(workdir::CommandOutputRequest {
|
||||
handle: CommandHandle("external".into()),
|
||||
cursor: 0,
|
||||
limit: 3,
|
||||
wait: true,
|
||||
})
|
||||
.await
|
||||
})
|
||||
};
|
||||
let (operation, response) = receive(&mut receiver).await;
|
||||
assert!(
|
||||
matches!(operation, WorkdirSessionOperation::CommandOutput(ref request) if request.handle.0 == "provider" && request.cursor == 0)
|
||||
);
|
||||
response
|
||||
.send(Ok(WorkdirSessionOperationResult::CommandOutput(page(
|
||||
0, b"end", true,
|
||||
))))
|
||||
.unwrap();
|
||||
assert_eq!(reader.await.unwrap().unwrap(), page(0, b"end", true));
|
||||
}
|
||||
assert_eq!(session.active_commands.lock().await.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_registry_retention_is_bounded_and_excludes_active_commands() {
|
||||
let (source, _) = command_fixture();
|
||||
let worker = RuntimeWorkerRef::new("runtime", "worker");
|
||||
let mut registry = WorkdirSessionRegistry::default();
|
||||
let active = registry.register_command(
|
||||
worker.clone(),
|
||||
"checkout".into(),
|
||||
source.clone(),
|
||||
CommandHandle("active".into()),
|
||||
);
|
||||
let mut last = None;
|
||||
for index in 0..MAX_RETAINED_TERMINAL_COMMANDS + 4 {
|
||||
let handle = registry.register_command(
|
||||
worker.clone(),
|
||||
"checkout".into(),
|
||||
source.clone(),
|
||||
CommandHandle(format!("terminal-{index}")),
|
||||
);
|
||||
registry.mark_terminal(&worker, &handle);
|
||||
last = Some(handle);
|
||||
}
|
||||
assert_eq!(registry.active_command_count(&worker), 1);
|
||||
assert_eq!(registry.commands.len(), MAX_RETAINED_TERMINAL_COMMANDS + 1);
|
||||
assert!(registry.command(&worker, "checkout", &active).is_some());
|
||||
assert!(
|
||||
registry
|
||||
.command(&worker, "checkout", &last.unwrap())
|
||||
.is_some()
|
||||
);
|
||||
}
|
||||
@@ -777,6 +777,7 @@ async fn drive_read_entrypoints_use_current_grants_and_cannot_fallback_to_browse
|
||||
configure_runtime_request_auth(&mut api, &identity, &workers[0].runtime_id);
|
||||
let routes = [
|
||||
format!("{}/root", base()),
|
||||
format!("{}/capabilities", base()),
|
||||
format!(
|
||||
"{}/metadata?entry_workspace_id={TEST_WORKSPACE_ID}&id={id}",
|
||||
base()
|
||||
@@ -1007,3 +1008,106 @@ async fn drive_download_stream_expires_instead_of_mixing_last_mutation_ids() {
|
||||
"next chunk must expire, not deliver generation b"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drive_capabilities_expose_only_authenticated_current_caller_ceiling() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let mut api = test_api(temp.path()).await;
|
||||
let owner = seed_test_api_token(api.store.as_ref(), TEST_WORKSPACE_ID);
|
||||
let (_runtime, workers) = workers(&api).await;
|
||||
let identity = RuntimeIdentityMaterial::generate(&workers[0].runtime_id).unwrap();
|
||||
configure_runtime_request_auth(&mut api, &identity, &workers[0].runtime_id);
|
||||
let path = format!("{}/capabilities", base());
|
||||
let unsigned = build_router(api.clone())
|
||||
.oneshot(Request::builder().uri(&path).body(Body::empty()).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
unsigned.status(),
|
||||
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
|
||||
));
|
||||
signed(
|
||||
&api,
|
||||
&identity,
|
||||
None,
|
||||
"GET",
|
||||
&path,
|
||||
None,
|
||||
StatusCode::FORBIDDEN,
|
||||
)
|
||||
.await;
|
||||
let mut grants = vec![];
|
||||
for (worker, access, write) in [
|
||||
(&workers[0], "read_only", false),
|
||||
(&workers[1], "read_write", true),
|
||||
] {
|
||||
grants.push(user_call(&api,&owner,"POST",&format!("{}/grants",base()),Some(json!({"runtime_id":worker.runtime_id,"worker_id":worker.worker_id,"access":access})),StatusCode::OK).await);
|
||||
assert_eq!(
|
||||
signed(
|
||||
&api,
|
||||
&identity,
|
||||
Some(&worker.worker_id),
|
||||
"GET",
|
||||
&path,
|
||||
None,
|
||||
StatusCode::OK
|
||||
)
|
||||
.await,
|
||||
json!({"write":write})
|
||||
);
|
||||
}
|
||||
// Altering a signed source cannot acquire the other Worker's larger ceiling.
|
||||
let mut forged =
|
||||
runtime_source_request(&identity, Some(&workers[0].worker_id), "GET", &path, vec![]);
|
||||
forged
|
||||
.headers_mut()
|
||||
.insert("x-yoi-worker-id", workers[1].worker_id.parse().unwrap());
|
||||
let response = build_router(api.clone()).oneshot(forged).await.unwrap();
|
||||
assert!(matches!(
|
||||
response.status(),
|
||||
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
|
||||
));
|
||||
let foreign = path.replace(&format!("/w/{TEST_WORKSPACE_ID}/"), "/w/other-workspace/");
|
||||
let response = build_router(api.clone())
|
||||
.oneshot(runtime_source_request(
|
||||
&identity,
|
||||
Some(&workers[1].worker_id),
|
||||
"GET",
|
||||
&foreign,
|
||||
vec![],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!response.status().is_success());
|
||||
signed(
|
||||
&api,
|
||||
&identity,
|
||||
Some("nonexistent-worker"),
|
||||
"GET",
|
||||
&path,
|
||||
None,
|
||||
StatusCode::UNAUTHORIZED,
|
||||
)
|
||||
.await;
|
||||
for (worker, grant) in workers.iter().zip(&grants) {
|
||||
user_call(
|
||||
&api,
|
||||
&owner,
|
||||
"DELETE",
|
||||
&format!("{}/grants/{}", base(), grant["grant_id"].as_str().unwrap()),
|
||||
None,
|
||||
StatusCode::OK,
|
||||
)
|
||||
.await;
|
||||
signed(
|
||||
&api,
|
||||
&identity,
|
||||
Some(&worker.worker_id),
|
||||
"GET",
|
||||
&path,
|
||||
None,
|
||||
StatusCode::FORBIDDEN,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -517,7 +517,7 @@ async fn user_uploaded_markdown_roundtrips_through_workdirless_tools_and_native_
|
||||
read["entry"]["metadata"]["last_mutation_id"],
|
||||
uploaded["entry"]["last_mutation_id"]
|
||||
);
|
||||
value(native_tool(&tools, "Invoke", json!({"path":path,"interface":reference,"operation":"write","arguments":{"content":updated_text}})).await.unwrap())
|
||||
value(native_tool(&tools, "Invoke", json!({"path":path,"interface":file_write_reference(&inspected),"operation":"write","arguments":{"content":updated_text}})).await.unwrap())["metadata"].clone()
|
||||
} else {
|
||||
let read = value(
|
||||
tool(&f.features[0], "DriveRead", json!({"entry":id}))
|
||||
@@ -733,6 +733,19 @@ async fn worker_removal_preserves_drive_reference_for_explicitly_granted_replace
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(f.api.runtime.worker(old).is_err());
|
||||
let client = f.clients[0].clone();
|
||||
let capability = tokio::task::spawn_blocking(move || {
|
||||
client.request(
|
||||
WorkspaceRequestMethod::Get,
|
||||
format!("{}/capabilities", base()),
|
||||
vec![],
|
||||
65536,
|
||||
)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(!(200..300).contains(&capability.status));
|
||||
assert!(
|
||||
f.api
|
||||
.store
|
||||
@@ -1571,11 +1584,30 @@ impl agen::llm_client::LlmClient for NoLlm {
|
||||
panic!("native tool fixture must not call an LLM")
|
||||
}
|
||||
}
|
||||
fn file_write_reference(observed: &Value) -> Value {
|
||||
observed["interfaces"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.find(|i| i["reference"]["name"] == "yoi.file/write-text/v1")
|
||||
.expect("write capability")["reference"]
|
||||
.clone()
|
||||
}
|
||||
fn native_tools(
|
||||
feature: &DriveFeature,
|
||||
) -> std::collections::BTreeMap<String, Arc<dyn agen::tool::Tool>> {
|
||||
native_tools_with_checkouts(feature, None)
|
||||
}
|
||||
fn native_tools_with_checkouts(
|
||||
feature: &DriveFeature,
|
||||
router: Option<Arc<WorkdirSessionRouter>>,
|
||||
) -> std::collections::BTreeMap<String, Arc<dyn agen::tool::Tool>> {
|
||||
let mut mounts = worker::wip::WipMountRegistry::new();
|
||||
worker::feature::builtin::drive::wip::mount_drive_wip(&mut mounts, feature).unwrap();
|
||||
if let Some(router) = router {
|
||||
worker::checkout::mount_checkouts(&mut mounts, router, tools::Tracker::new(), None)
|
||||
.unwrap();
|
||||
}
|
||||
let mut engine = agen::Engine::<_, agen::state::Mutable, ()>::new_annotated(NoLlm);
|
||||
engine.register_tools(
|
||||
feature
|
||||
@@ -1651,7 +1683,14 @@ async fn native_tree_does_not_index_entries_but_direct_inspect_and_invoke_use_cu
|
||||
);
|
||||
assert_eq!(read["text"], "native content");
|
||||
let denied_write = native_tool(&tools, "Invoke", json!({"path":path,"interface":reference,"operation":"write","arguments":{"content":"lost"}})).await;
|
||||
denied(denied_write);
|
||||
assert!(denied_write.is_err());
|
||||
assert!(
|
||||
!inspected["interfaces"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.any(|i| i["reference"]["name"] == "yoi.file/write-text/v1")
|
||||
);
|
||||
f.revoke(1).await;
|
||||
assert!(
|
||||
native_tool(
|
||||
@@ -1705,12 +1744,12 @@ async fn native_create_and_write_use_client_managed_committed_requests_and_rejec
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let updated = value(native_tool(&first, "Invoke", json!({"path":path,"interface":observed_first["interfaces"][0]["reference"],"operation":"write","arguments":{"content":"committed"}})).await.unwrap());
|
||||
let updated = value(native_tool(&first, "Invoke", json!({"path":path,"interface":file_write_reference(&observed_first),"operation":"write","arguments":{"content":"committed"}})).await.unwrap())["metadata"].clone();
|
||||
assert_ne!(
|
||||
updated["entry"]["metadata"]["last_mutation_id"],
|
||||
created["entry"]["metadata"]["last_mutation_id"]
|
||||
);
|
||||
let error = native_tool(&second, "Invoke", json!({"path":path,"interface":observed_second["interfaces"][0]["reference"],"operation":"write","arguments":{"content":"lost"}})).await.unwrap_err();
|
||||
let error = native_tool(&second, "Invoke", json!({"path":path,"interface":file_write_reference(&observed_second),"operation":"write","arguments":{"content":"lost"}})).await.unwrap_err();
|
||||
assert!(
|
||||
error.to_string().to_lowercase().contains("validator"),
|
||||
"stale native mutation must report validator conflict: {error:?}"
|
||||
@@ -1833,3 +1872,143 @@ async fn image_attachment_retains_original_bytes_after_drive_update_and_revocati
|
||||
"history attachment must not fetch latest Drive content"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn common_file_inspect_invoke_crosses_drive_and_checkout_without_transferring_authority() {
|
||||
let f = Fixture::new(["read_write", "read_only"]).await;
|
||||
let created = f.create("common.md", "first\n雪 needle\nlast\n").await;
|
||||
let path = created["entry"]["path"].as_str().unwrap();
|
||||
let checkout = f._temp.path().join("checkout");
|
||||
fs::create_dir(&checkout).unwrap();
|
||||
fs::write(checkout.join("file"), "first\n雪 needle\nlast\n").unwrap();
|
||||
let router = Arc::new(WorkdirSessionRouter::new());
|
||||
router
|
||||
.attach(
|
||||
workdir::WorkdirAttachmentAlias::new("main").unwrap(),
|
||||
Arc::new(workdir::LocalWorkdirSession::materialized(
|
||||
checkout.clone(),
|
||||
checkout.clone(),
|
||||
manifest::SharedScope::new(manifest::Scope::writable(&checkout).unwrap()),
|
||||
workdir::WorkdirSessionCapabilities::READ_WRITE,
|
||||
)),
|
||||
)
|
||||
.unwrap();
|
||||
let tools = native_tools_with_checkouts(&f.features[0], Some(router));
|
||||
let drive = value(
|
||||
native_tool(&tools, "Inspect", json!({"path":path}))
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let local = value(
|
||||
native_tool(&tools, "Inspect", json!({"path":"/checkouts/main/file"}))
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
for name in [
|
||||
"yoi.file/read-lines/v1",
|
||||
"yoi.file/search-text/v1",
|
||||
"yoi.file/edit-text/v1",
|
||||
"yoi.file/write-text/v1",
|
||||
] {
|
||||
let selected = drive["interfaces"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.find(|i| i["reference"]["name"] == name)
|
||||
.unwrap();
|
||||
let reference = &selected["reference"];
|
||||
assert_eq!(reference["scope"], "/");
|
||||
assert!(
|
||||
local["interfaces"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.any(|i| i["reference"] == *reference)
|
||||
);
|
||||
let signature = value(
|
||||
native_tool(&tools, "Inspect", json!({"path":selected["path"]}))
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
assert!(
|
||||
signature["interface_signature"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains(name)
|
||||
);
|
||||
let (operation, arguments) = match name {
|
||||
"yoi.file/read-lines/v1" => ("read", json!({"offset":1,"limit":1})),
|
||||
"yoi.file/search-text/v1" => ("search", json!({"pattern":"needle","limit":2})),
|
||||
"yoi.file/edit-text/v1" => {
|
||||
("edit", json!({"old_string":"needle","new_string":"found"}))
|
||||
}
|
||||
_ => ("write", json!({"content":"saved\n"})),
|
||||
};
|
||||
let mut results = vec![];
|
||||
for target in [path, "/checkouts/main/file"] {
|
||||
results.push(value(native_tool(&tools,"Invoke",json!({"path":target,"interface":reference,"operation":operation,"arguments":arguments})).await.unwrap()));
|
||||
}
|
||||
if matches!(operation, "read" | "search") {
|
||||
assert_eq!(results[0], results[1]);
|
||||
} else {
|
||||
assert!(results.iter().all(|v| v.get("metadata").is_some()));
|
||||
}
|
||||
}
|
||||
assert_eq!(
|
||||
fs::read_to_string(checkout.join("file")).unwrap(),
|
||||
"saved\n"
|
||||
);
|
||||
assert_eq!(
|
||||
value(
|
||||
tool(
|
||||
&f.features[1],
|
||||
"DriveRead",
|
||||
json!({"entry":entry(&created)})
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
)["text"],
|
||||
"saved\n"
|
||||
);
|
||||
// A capability query is not a grant: drop the writer to readonly after the common write was cached.
|
||||
f.revoke(0).await;
|
||||
owner_call(&f.api,&f.owner,"POST",&format!("{}/grants",base()),Some(json!({"runtime_id":f.clients[0].worker.runtime_id,"worker_id":f.clients[0].worker.worker_id,"access":"read_only"}))).await;
|
||||
assert!(native_tool(&tools,"Invoke",json!({"path":path,"interface":file_write_reference(&drive),"operation":"write","arguments":{"content":"lost"}})).await.is_err());
|
||||
let refreshed = value(
|
||||
native_tool(&tools, "Inspect", json!({"path":path,"refresh":true}))
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
assert!(
|
||||
!refreshed["interfaces"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.any(|i| i["reference"]["name"] == "yoi.file/write-text/v1")
|
||||
);
|
||||
let common_read = json!({"scope":"/","name":"yoi.file/read-lines/v1"});
|
||||
assert_eq!(
|
||||
value(
|
||||
native_tool(
|
||||
&tools,
|
||||
"Invoke",
|
||||
json!({"path":path,"interface":common_read,"operation":"read","arguments":{}})
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
)["content"],
|
||||
" 1\tsaved\n"
|
||||
);
|
||||
assert_eq!(
|
||||
value(
|
||||
tool(
|
||||
&f.features[1],
|
||||
"DriveRead",
|
||||
json!({"entry":entry(&created)})
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
)["text"],
|
||||
"saved\n"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -256,6 +256,7 @@ fn read_result() -> WorkdirSessionOperationResult {
|
||||
total_lines: 1,
|
||||
content_hash: sha2::Sha256::digest(b"visible\n").into(),
|
||||
truncated: false,
|
||||
byte_truncated: false,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -466,7 +466,6 @@ impl AttentionWorkerFactory {
|
||||
worker.with_notification_coalesce_delay(std::time::Duration::from_millis(10)),
|
||||
worker::WorkerBootstrapLayout::RuntimeManagedRun {
|
||||
run_dir: root.join(uuid::Uuid::now_v7().to_string()),
|
||||
bash_output_dir: root.join("bash-output"),
|
||||
},
|
||||
worker::WorkerControllerTransport::InProcess,
|
||||
)
|
||||
|
||||
@@ -799,7 +799,6 @@ impl DeliveryWorkerFactory {
|
||||
worker::WorkerFilesystemAuthority::None,
|
||||
worker::WorkerBootstrapLayout::RuntimeManagedRun {
|
||||
run_dir: root.join(uuid::Uuid::now_v7().to_string()),
|
||||
bash_output_dir: root.join("bash-output"),
|
||||
},
|
||||
worker::WorkerControllerTransport::InProcess,
|
||||
)
|
||||
|
||||
@@ -238,6 +238,25 @@ pub(super) async fn root(
|
||||
let drive = drive_grants::authorized_drive(api, context, workspace, false).await?;
|
||||
blocking(false, move || drive.root()).await.map(public_node)
|
||||
}
|
||||
/// Query only the authenticated caller's current ceiling using the operation authorizer.
|
||||
/// This is admission/publication evidence, never authority retained for a later mutation.
|
||||
pub(super) async fn capabilities(
|
||||
api: &WorkspaceApi,
|
||||
context: &ServerRequestContext,
|
||||
workspace: &str,
|
||||
) -> DriveResult<DriveCapabilities> {
|
||||
let drive = drive_grants::authorized_drive(api, context, workspace, false).await?;
|
||||
let write = blocking(false, move || match drive.check_authorized(true) {
|
||||
Ok(()) => Ok(true),
|
||||
Err(workspace_drive::Error::Denied) => {
|
||||
drive.check_authorized(false)?;
|
||||
Ok(false)
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
})
|
||||
.await?;
|
||||
Ok(DriveCapabilities { write })
|
||||
}
|
||||
pub(super) async fn metadata(
|
||||
api: &WorkspaceApi,
|
||||
context: &ServerRequestContext,
|
||||
@@ -581,6 +600,9 @@ pub(super) fn router(service: Arc<ServerApiContractService>) -> Router {
|
||||
.merge(streaming_router(service.clone()))
|
||||
.merge(server_api::server_api_axum::drive_root(service.clone()))
|
||||
.merge(server_api::server_api_axum::drive_metadata(service.clone()))
|
||||
.merge(server_api::server_api_axum::drive_capabilities(
|
||||
service.clone(),
|
||||
))
|
||||
.merge(server_api::server_api_axum::drive_list(service.clone()))
|
||||
.merge(server_api::server_api_axum::drive_search(service.clone()))
|
||||
.merge(server_api::server_api_axum::drive_read_text(
|
||||
|
||||
@@ -187,7 +187,7 @@ async fn config_call(
|
||||
runtime
|
||||
.invoke(
|
||||
path.into(),
|
||||
config_interface(path),
|
||||
if matches!(operation,"edit"|"write") {wip_protocol::InterfaceReference {scope:"/".into(),name:format!("yoi.file/{operation}-text/v1")}} else {config_interface(path)},
|
||||
operation.into(),
|
||||
args,
|
||||
agen::tool::ToolExecutionContext::direct(),
|
||||
|
||||
Reference in New Issue
Block a user