From cc2f8b3422d0ac53b5d21b79313fdf1c40cb7899 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 11:07:59 +0900 Subject: [PATCH 01/11] feat: add opt-in early tool call dispatch --- crates/agen/src/engine.rs | 751 +++++++++++++++++- crates/agen/src/lib.rs | 2 +- crates/agen/src/llm_client/client.rs | 19 + .../scheme/anthropic/scheme_impl.rs | 5 + crates/agen/src/llm_client/scheme/mod.rs | 33 +- .../scheme/openai_responses/scheme_impl.rs | 6 +- crates/agen/src/llm_client/transport.rs | 4 + crates/agen/src/timeline/timeline.rs | 79 +- .../agen/src/timeline/tool_call_collector.rs | 51 +- crates/agen/tests/parallel_execution_test.rs | 472 ++++++++++- crates/manifest/src/config.rs | 37 + crates/manifest/src/lib.rs | 17 + crates/worker/src/spawn/tool.rs | 8 +- crates/worker/src/worker.rs | 14 + docs/manifest.toml | 8 + 15 files changed, 1433 insertions(+), 73 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 1fa4fafd0..2ee42feaa 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -1,7 +1,8 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::{future::Future, marker::PhantomData, pin::Pin, sync::Arc, time::Instant}; use futures::StreamExt; +use futures::stream::FuturesUnordered; use serde_json::{Value, json}; use tokio::sync::mpsc; use tokio::time::Instant as TokioInstant; @@ -25,7 +26,7 @@ use crate::{ llm_client::{ ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream, ToolDefinition, error::is_retryable, event::Event, retry::RetryPolicy, - transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT, types::parse_tool_arguments, + transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT, }, state::{EngineState, Locked, Mutable}, timeline::event::{ErrorEvent, StatusEvent, UsageEvent}, @@ -78,11 +79,20 @@ pub enum ToolRegistryError { DuplicateName(String), } +/// When provider-confirmed tool calls may begin execution. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum ToolCallDispatchMode { + /// Wait for the complete assistant response before starting any tools. + #[default] + AfterResponse, + /// Start calls at trustworthy provider per-call completion boundaries. + OnToolCallComplete, +} + /// Engine configuration #[derive(Debug, Clone, Default)] pub struct EngineConfig { - // Reserved for future extensions (currently empty) - _private: (), + pub tool_call_dispatch: ToolCallDispatchMode, } /// Project terminal tool outputs into the assistant's original ToolCall order. @@ -196,6 +206,7 @@ pub struct EngineRunOutput { } /// Internal: tool execution result +#[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ToolExecutionResult { Completed, Paused, @@ -207,6 +218,56 @@ struct ToolExecutionAttempt { terminal: bool, } +type ToolTerminalTask = Pin + Send>>; +type ToolCallExecutionInfo = ( + ToolCall, + crate::tool::ToolMeta, + Arc, + ToolExecutionContext, +); + +/// One response-scoped execution batch whose calls are admitted as their +/// provider-confirmed blocks complete. +struct EarlyToolExecutionBatch { + batch_id: String, + next_call_index: usize, + calls: Vec, + seen_call_ids: HashSet, + call_info_map: HashMap, + started_calls: Vec<(String, String)>, + attempt_fence: ToolExecutionAttemptFence, + futures: FuturesUnordered, + execution_handles: HashMap, + synthetic_results: VecDeque, + terminal_call_ids: HashSet, + dispatched_any: bool, +} + +impl EarlyToolExecutionBatch { + fn new(batch_id: String) -> Self { + Self { + batch_id, + next_call_index: 0, + calls: Vec::new(), + seen_call_ids: HashSet::new(), + call_info_map: HashMap::new(), + started_calls: Vec::new(), + attempt_fence: ToolExecutionAttemptFence::default(), + futures: FuturesUnordered::new(), + execution_handles: HashMap::new(), + synthetic_results: VecDeque::new(), + terminal_call_ids: HashSet::new(), + dispatched_any: false, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum EarlyToolAdmission { + Continue, + Pause, +} + /// Per-batch compare-and-set fence for terminal ToolResult commits. /// /// A completion may commit only when its attempt id still matches the active @@ -312,6 +373,11 @@ enum StreamCompletion { Interrupted { reason: String }, } +struct StreamResponseOutput { + completion: StreamCompletion, + early_tools: Option, +} + pub struct Engine { /// LLM client client: C, @@ -362,6 +428,8 @@ pub struct Engine { tool_execution_batch_count: usize, /// Maximum number of AgentTurns (None = unlimited) max_turns: Option, + /// Selected tool-call execution start boundary. + tool_call_dispatch_mode: ToolCallDispatchMode, /// Caller-selected policy for interrupting started provider operations. tool_execution_policy: ToolExecutionPolicy, /// AgentTurn-start callbacks (1:1 with LlmCall today) @@ -1184,7 +1252,10 @@ impl Engine { } = item { if !answered_call_ids.contains(call_id) { - let input = parse_tool_arguments(arguments); + let input = serde_json::from_str::(arguments) + .ok() + .filter(Value::is_object) + .unwrap_or(Value::Null); pending_calls.push(ToolCall { id: call_id.clone(), name: name.clone(), @@ -1201,6 +1272,342 @@ impl Engine { } } + async fn admit_early_tool_calls( + &mut self, + history: &History, + batch: &mut EarlyToolExecutionBatch, + tool_calls: Vec, + ) -> Result { + for mut tool_call in tool_calls { + if !batch.seen_call_ids.insert(tool_call.id.clone()) { + warn!(call_id = %tool_call.id, "Ignoring duplicate completed tool-call event"); + continue; + } + let call_index = batch.next_call_index; + batch.next_call_index += 1; + batch.calls.push(tool_call.clone()); + let expected_tool_use_id = tool_call.id.clone(); + let context = ToolExecutionContext::new(&tool_call.id, &batch.batch_id, call_index); + + if !tool_call.input.is_object() { + batch.synthetic_results.push_back(ToolResult::error( + &tool_call.id, + "Tool arguments must be one complete JSON object", + )); + batch.dispatched_any = true; + continue; + } + + let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) else { + let attempt_id = context.batch_id.clone(); + batch + .attempt_fence + .register(tool_call.id.clone(), attempt_id.clone()); + batch + .started_calls + .push((tool_call.id.clone(), attempt_id.clone())); + let call_id = tool_call.id.clone(); + let name = tool_call.name.clone(); + batch.futures.push(Box::pin(async move { + ( + attempt_id, + ToolResult::error(&call_id, format!("Tool not found: {name}")), + ) + })); + batch.dispatched_any = true; + continue; + }; + + let invocation = self.interceptor_invocation( + InterceptorPhase::PreToolCall, + Some(self.turn_count), + Some(InterceptorCallId::Tool(expected_tool_use_id.clone())), + call_index, + ); + let mut info = ToolCallInfo { + invocation, + history: history.entries(), + call: tool_call.clone(), + meta, + tool, + context, + }; + let action = self + .interceptor + .pre_tool_call(&mut info) + .await + .map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPhase::PreToolCall, + error, + )) + })?; + if info.call.id != expected_tool_use_id { + return Err(InterceptorFailure::new( + InterceptorPhase::PreToolCall, + InterceptorError::new( + InterceptorErrorCategory::ContractViolation, + "pre-tool interceptor changed immutable tool call identity", + ), + ) + .into()); + } + tool_call = info.call; + let mut context = info.context; + context.call_id = tool_call.id.clone(); + batch.call_info_map.insert( + tool_call.id.clone(), + ( + tool_call.clone(), + info.meta.clone(), + info.tool.clone(), + context.clone(), + ), + ); + + match action { + PreToolAction::Continue => { + let attempt_id = context.batch_id.clone(); + batch + .attempt_fence + .register(tool_call.id.clone(), attempt_id.clone()); + batch + .started_calls + .push((tool_call.id.clone(), attempt_id.clone())); + let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default(); + let call_id = tool_call.id.clone(); + let (handle, terminal) = + ToolExecutionHandle::start(info.tool, input_json, context); + batch.execution_handles.insert(call_id.clone(), handle); + batch.futures.push(Box::pin(async move { + let result = match terminal.await { + ToolExecutionTerminal::Confirmed(Ok(output)) => { + ToolResult::from_output(&call_id, output) + } + ToolExecutionTerminal::Confirmed(Err(ToolError::Cancelled(output))) => { + ToolResult::from_output_with_disposition( + &call_id, + output, + ToolResultDisposition::Cancelled, + ) + } + ToolExecutionTerminal::Confirmed(Err(ToolError::Interrupted( + output, + ))) => ToolResult::from_output_with_disposition( + &call_id, + output, + ToolResultDisposition::Interrupted, + ), + ToolExecutionTerminal::Confirmed(Err(error)) => { + ToolResult::error(&call_id, error.to_string()) + } + ToolExecutionTerminal::OutcomeUnknown => { + ToolResult::outcome_unknown(&call_id) + } + }; + (attempt_id, result) + })); + batch.dispatched_any = true; + } + PreToolAction::SyntheticResult(result) => { + if result.tool_use_id != expected_tool_use_id { + return Err(InterceptorFailure::new( + InterceptorPhase::PreToolCall, + InterceptorError::new( + InterceptorErrorCategory::ContractViolation, + "synthetic tool result changed immutable tool call identity", + ), + ) + .into()); + } + batch.synthetic_results.push_back(result); + batch.dispatched_any = true; + } + PreToolAction::Skip => {} + PreToolAction::Abort(reason) => return Err(EngineError::Aborted(reason)), + PreToolAction::Pause => return Ok(EarlyToolAdmission::Pause), + } + } + Ok(EarlyToolAdmission::Continue) + } + + async fn cancel_and_terminalize_early_batch( + &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + batch: &mut EarlyToolExecutionBatch, + ) -> Option { + use futures::StreamExt as _; + + let cancellation_request_deadline = + TokioInstant::now() + self.tool_execution_policy.cancellation_request_timeout; + let cancellation_requests = batch + .execution_handles + .iter() + .filter(|(call_id, _)| !batch.terminal_call_ids.contains(*call_id)) + .map(|(call_id, handle)| { + let call_id = call_id.clone(); + let handle = handle.clone(); + async move { + ( + call_id, + handle.cancel_before(cancellation_request_deadline).await, + ) + } + }); + let cancellation_requests: FuturesUnordered<_> = cancellation_requests.collect(); + for (call_id, result) in cancellation_requests.collect::>().await { + if let Err(error) = result { + warn!(%call_id, error = %error, "Tool cooperative cancellation request failed"); + } + } + + let mut batch_error = None; + let deadline = + TokioInstant::now() + self.tool_execution_policy.terminal_confirmation_timeout; + while !batch.futures.is_empty() { + tokio::select! { + biased; + result = batch.futures.next() => { + let (attempt_id, result) = result.expect("non-empty FuturesUnordered"); + if let Err(error) = self.finalize_and_commit_tool_result( + history, + annotate, + result, + Some(&attempt_id), + &batch.call_info_map, + &mut batch.attempt_fence, + &mut batch.terminal_call_ids, + ).await && batch_error.is_none() { + batch_error = Some(error); + } + } + _ = tokio::time::sleep_until(deadline) => break, + } + } + for (call_id, attempt_id) in &batch.started_calls { + if !batch.attempt_fence.is_terminal(call_id) { + if let Some(handle) = batch.execution_handles.get(call_id) { + handle.force_close(); + } + if let Err(error) = self + .finalize_and_commit_tool_result( + history, + annotate, + ToolResult::outcome_unknown(call_id), + Some(attempt_id), + &batch.call_info_map, + &mut batch.attempt_fence, + &mut batch.terminal_call_ids, + ) + .await + && batch_error.is_none() + { + batch_error = Some(error); + } + } + } + batch_error + } + + async fn drain_early_tool_batch( + &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + mut batch: EarlyToolExecutionBatch, + force_cancel: bool, + initial_pause_requested: bool, + ) -> Result { + while let Some(result) = batch.synthetic_results.pop_front() { + if let Err(error) = self + .finalize_and_commit_tool_result( + history, + annotate, + result, + None, + &batch.call_info_map, + &mut batch.attempt_fence, + &mut batch.terminal_call_ids, + ) + .await + { + let cleanup_error = self + .cancel_and_terminalize_early_batch(history, annotate, &mut batch) + .await; + return Err(cleanup_error.unwrap_or(error)); + } + } + + if force_cancel { + if let Some(error) = self + .cancel_and_terminalize_early_batch(history, annotate, &mut batch) + .await + { + return Err(error); + } + return Ok(ToolExecutionResult::Completed); + } + + let mut pause_requested = initial_pause_requested; + let mut pause_deadline = initial_pause_requested + .then(|| TokioInstant::now() + self.tool_execution_policy.pause_safe_boundary_timeout); + while !batch.futures.is_empty() { + tokio::select! { + biased; + result = batch.futures.next() => { + let (attempt_id, result) = result.expect("non-empty FuturesUnordered"); + if let Err(error) = self.finalize_and_commit_tool_result( + history, + annotate, + result, + Some(&attempt_id), + &batch.call_info_map, + &mut batch.attempt_fence, + &mut batch.terminal_call_ids, + ).await { + let cleanup_error = self + .cancel_and_terminalize_early_batch(history, annotate, &mut batch) + .await; + return Err(cleanup_error.unwrap_or(error)); + } + } + pause = self.pause_rx.recv(), if !pause_requested => { + if pause.is_some() { + pause_requested = true; + pause_deadline = Some( + TokioInstant::now() + + self.tool_execution_policy.pause_safe_boundary_timeout, + ); + } + } + _ = tokio::time::sleep_until(pause_deadline.unwrap_or_else(TokioInstant::now)), if pause_deadline.is_some() => { + if let Some(error) = self + .cancel_and_terminalize_early_batch(history, annotate, &mut batch) + .await { + return Err(error); + } + return Ok(ToolExecutionResult::Paused); + } + cancel = self.cancel_rx.recv() => { + if cancel.is_some() { + info!("Tool execution cancellation requested"); + } + if let Some(error) = self + .cancel_and_terminalize_early_batch(history, annotate, &mut batch) + .await { + return Err(error); + } + return Err(EngineError::Cancelled); + } + } + } + Ok(if pause_requested { + ToolExecutionResult::Paused + } else { + ToolExecutionResult::Completed + }) + } + /// Execute tools in parallel /// /// After running pre_tool_call hooks on all tools, @@ -1231,6 +1638,13 @@ impl Engine { for (call_index, mut tool_call) in tool_calls.into_iter().enumerate() { let expected_tool_use_id = tool_call.id.clone(); let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index); + if !tool_call.input.is_object() { + synthetic_results.push(ToolResult::error( + &tool_call.id, + "Tool arguments must be one complete JSON object", + )); + continue; + } if let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) { let invocation = self.interceptor_invocation( InterceptorPhase::PreToolCall, @@ -1934,8 +2348,30 @@ impl Engine { self.request_trace_payload(&request), ); let request = self.attach_transport_trace(request, current_turn, current_llm_call); - let stream_outcome = self - .stream_response(request, current_turn, current_llm_call) + let early_dispatch = self.tool_call_dispatch_mode + == ToolCallDispatchMode::OnToolCallComplete + && self.client.tool_call_completion_support() + == crate::llm_client::ToolCallCompletionSupport::PerBlock; + if self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete + && !early_dispatch + { + self.emit_warning( + "tool-call early dispatch deferred until response completion: provider does not expose a trustworthy per-call completion boundary", + ); + } + let assistant_start = history.len(); + let StreamResponseOutput { + completion: stream_outcome, + mut early_tools, + } = self + .stream_response( + request, + current_turn, + current_llm_call, + history, + annotate, + early_dispatch, + ) .await?; for cb in &self.llm_call_end_cbs { @@ -1944,6 +2380,47 @@ impl Engine { self.llm_call_count += 1; if let StreamCompletion::Interrupted { reason } = stream_outcome { + if early_tools + .as_ref() + .is_some_and(|batch| batch.dispatched_any) + { + // Once a side effect may have started, automatic stream + // continuation would allow the model to regenerate the same + // operation under a fresh call id. Close every owned attempt + // and stop this run instead. + self.timeline.abort_current_block(); + self.timeline.flush_usage(); + let reasoning_items = self.thinking_block_collector.take_collected(); + let text_blocks = self.text_block_collector.take_collected(); + let _partial_tool_calls = self.tool_call_collector.take_collected(); + let assistant_items = + self.build_assistant_items(&reasoning_items, &text_blocks, &[]); + if !assistant_items.is_empty() + && let Err(error) = + self.append_history_items(history, assistant_items, annotate) + { + if let Some(batch) = early_tools.take() + && let Err(cleanup_error) = self + .drain_early_tool_batch(history, annotate, batch, true, false) + .await + { + return Err(cleanup_error); + } + return Err(error); + } + if let Some(batch) = early_tools.take() { + self.drain_early_tool_batch(history, annotate, batch, true, false) + .await?; + } + return Err(EngineError::Client(ClientError::Api { + status: None, + code: Some("early_tool_stream_interrupted".to_string()), + message: format!( + "LLM stream interrupted after early tool execution started: {reason}" + ), + retry_after: None, + })); + } stream_continuations += 1; if stream_continuations > MAX_STREAM_CONTINUATIONS { return Err(EngineError::Client(ClientError::Api { @@ -1989,11 +2466,61 @@ impl Engine { // `append_history_items` so observers see each item as it lands. let reasoning_items = self.thinking_block_collector.take_collected(); let text_blocks = self.text_block_collector.take_collected(); - let tool_calls = self.tool_call_collector.take_collected(); + let remaining_tool_calls = self.tool_call_collector.take_collected(); let assistant_items = - self.build_assistant_items(&reasoning_items, &text_blocks, &tool_calls); - let assistant_start = history.len(); - self.append_history_items(history, assistant_items, annotate)?; + self.build_assistant_items(&reasoning_items, &text_blocks, &remaining_tool_calls); + if let Err(error) = self.append_history_items(history, assistant_items, annotate) { + if let Some(batch) = early_tools.take() + && let Err(cleanup_error) = self + .drain_early_tool_batch(history, annotate, batch, true, false) + .await + { + return Err(cleanup_error); + } + return Err(error); + } + + if early_dispatch && !remaining_tool_calls.is_empty() { + if early_tools.is_none() { + let batch_id = format!("tool-batch-{}", self.tool_execution_batch_count); + self.tool_execution_batch_count += 1; + early_tools = Some(EarlyToolExecutionBatch::new(batch_id)); + } + let admission = self + .admit_early_tool_calls( + history, + early_tools.as_mut().expect("batch initialized"), + remaining_tool_calls.clone(), + ) + .await; + let admission = match admission { + Ok(admission) => admission, + Err(error) => { + if let Some(batch) = early_tools.take() + && let Err(cleanup_error) = self + .drain_early_tool_batch(history, annotate, batch, true, false) + .await + { + return Err(cleanup_error); + } + return Err(error); + } + }; + if admission == EarlyToolAdmission::Pause { + if let Some(batch) = early_tools.take() { + self.drain_early_tool_batch(history, annotate, batch, false, true) + .await?; + } + return Ok(EngineResult::Paused); + } + } + let mut tool_calls = early_tools + .as_ref() + .map(|batch| batch.calls.clone()) + .unwrap_or_default(); + if !early_dispatch { + tool_calls = remaining_tool_calls; + } let assistant_invocation = self.interceptor_invocation( InterceptorPhase::AssistantTurnEnd, @@ -2011,29 +2538,64 @@ impl Engine { history: history.entries(), tool_calls: &tool_calls, }) - .await - .map_err(|error| { - EngineError::from(InterceptorFailure::new( + .await; + let assistant_turn_action = match assistant_turn_action { + Ok(action) => action, + Err(error) => { + let error = EngineError::from(InterceptorFailure::new( InterceptorPhase::AssistantTurnEnd, error, - )) - })?; + )); + if let Some(batch) = early_tools.take() + && let Err(cleanup_error) = self + .drain_early_tool_batch(history, annotate, batch, true, false) + .await + { + return Err(cleanup_error); + } + return Err(error); + } + }; match assistant_turn_action { TurnEndAction::Finish if tool_calls.is_empty() => { return Ok(EngineResult::Finished); } TurnEndAction::Finish => {} TurnEndAction::ContinueWithMessages(additional) => { - self.append_history_items(history, additional, annotate)?; + if let Err(error) = self.append_history_items(history, additional, annotate) { + if let Some(batch) = early_tools.take() + && let Err(cleanup_error) = self + .drain_early_tool_batch(history, annotate, batch, true, false) + .await + { + return Err(cleanup_error); + } + return Err(error); + } if tool_calls.is_empty() { continue; } } TurnEndAction::Pause => { + if let Some(batch) = early_tools.take() { + self.drain_early_tool_batch(history, annotate, batch, false, true) + .await?; + } return Ok(EngineResult::Paused); } } + if let Some(batch) = early_tools.take() { + if self + .drain_early_tool_batch(history, annotate, batch, false, false) + .await? + == ToolExecutionResult::Paused + { + return Ok(EngineResult::Paused); + } + continue; + } + if let Some(result) = self .execute_and_commit_tools(history, annotate, tool_calls) .await? @@ -2224,7 +2786,10 @@ impl Engine { request: Request, turn: usize, llm_call: usize, - ) -> Result { + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + early_dispatch: bool, + ) -> Result { debug!( item_count = request.items.len(), tool_count = request.tools.len(), @@ -2233,6 +2798,7 @@ impl Engine { ); let mut stream = self.open_stream_with_retry(request, turn, llm_call).await?; + let mut early_tools: Option = None; let mut event_count: usize = 0; loop { @@ -2254,8 +2820,11 @@ impl Engine { Err(err) => { // 部分情報でも発火しておく（料金会計用） self.timeline.flush_usage(); - return Ok(StreamCompletion::Interrupted { - reason: err.to_string(), + return Ok(StreamResponseOutput { + completion: StreamCompletion::Interrupted { + reason: err.to_string(), + }, + early_tools, }); } }; @@ -2269,9 +2838,116 @@ impl Engine { } self.emit_stream_event(turn, llm_call, &event); self.timeline.dispatch(&event); + if early_dispatch { + let collected = self.tool_call_collector.take_collected(); + if !collected.is_empty() { + if early_tools.is_none() { + let batch_id = format!( + "tool-batch-{}", + self.tool_execution_batch_count + ); + self.tool_execution_batch_count += 1; + early_tools = + Some(EarlyToolExecutionBatch::new(batch_id)); + } + let fresh_calls = { + let batch = early_tools.as_ref().expect("batch initialized"); + collected + .into_iter() + .filter(|call| { + !batch.seen_call_ids.contains(&call.id) + }) + .collect::>() + }; + if !fresh_calls.is_empty() { + // Commit every completed assistant block observed so far, + // including the ToolCall itself, before provider ownership + // can start. This is the durable fail-closed gate. + let reasoning = + self.thinking_block_collector.take_collected(); + let text = self.text_block_collector.take_collected(); + let assistant_items = self.build_assistant_items( + &reasoning, + &text, + &fresh_calls, + ); + if let Err(error) = self.append_history_items( + history, + assistant_items, + annotate, + ) { + if let Some(mut batch) = early_tools.take() { + let _ = self + .cancel_and_terminalize_early_batch( + history, + annotate, + &mut batch, + ) + .await; + } + return Err(error); + } + let admission = { + let batch = + early_tools.as_mut().expect("batch initialized"); + self.admit_early_tool_calls( + history, + batch, + fresh_calls, + ) + .await + }; + match admission { + Ok(EarlyToolAdmission::Continue) => {} + Ok(EarlyToolAdmission::Pause) => { + self.timeline.abort_current_block(); + self.timeline.flush_usage(); + if let Some(batch) = early_tools.take() { + self.drain_early_tool_batch( + history, + annotate, + batch, + false, + true, + ) + .await?; + } + return Err(EngineError::PauseRequested); + } + Err(error) => { + self.timeline.abort_current_block(); + self.timeline.flush_usage(); + if let Some(mut batch) = early_tools.take() { + let cleanup_error = self + .cancel_and_terminalize_early_batch( + history, + annotate, + &mut batch, + ) + .await; + return Err(cleanup_error.unwrap_or(error)); + } + return Err(error); + } + } + } + } + } if let Event::Error(err) = &event { self.timeline.abort_current_block(); self.timeline.flush_usage(); + if let Some(mut batch) = early_tools.take() { + if let Some(error) = self + .cancel_and_terminalize_early_batch( + history, + annotate, + &mut batch, + ) + .await + { + return Err(error); + } + } return Err(EngineError::Client(ClientError::Api { status: None, code: err.code.clone(), @@ -2291,6 +2967,10 @@ impl Engine { // Preserve any UsageEvent already received for billing, but // do not treat this paused request as a normal completion. self.timeline.flush_usage(); + if let Some(batch) = early_tools.take() { + self.drain_early_tool_batch(history, annotate, batch, false, true) + .await?; + } return Err(EngineError::PauseRequested); } cancel = self.cancel_rx.recv() => { @@ -2299,6 +2979,14 @@ impl Engine { } self.timeline.abort_current_block(); self.timeline.flush_usage(); + if let Some(mut batch) = early_tools.take() { + if let Some(error) = self + .cancel_and_terminalize_early_batch(history, annotate, &mut batch) + .await + { + return Err(error); + } + } return Err(EngineError::Cancelled); } } @@ -2306,7 +2994,10 @@ impl Engine { // ストリーム完了時に集約済み Usage を 1 度だけ発火 self.timeline.flush_usage(); debug!(event_count = event_count, "Stream completed"); - Ok(StreamCompletion::Complete) + Ok(StreamResponseOutput { + completion: StreamCompletion::Complete, + early_tools, + }) } /// Execute tools and push results to history. @@ -2361,6 +3052,7 @@ impl Engine { llm_call_count: 0, tool_execution_batch_count: 0, max_turns: None, + tool_call_dispatch_mode: ToolCallDispatchMode::default(), tool_execution_policy: ToolExecutionPolicy::default(), turn_start_cbs: Vec::new(), turn_end_cbs: Vec::new(), @@ -2561,9 +3253,18 @@ impl Engine { self.max_turns = max_turns; } - /// Apply configuration (reserved for future extensions) - #[allow(dead_code)] - pub fn config(self, _config: EngineConfig) -> Self { + /// Select when provider-confirmed tool calls may begin execution. + pub fn set_tool_call_dispatch_mode(&mut self, mode: ToolCallDispatchMode) { + self.tool_call_dispatch_mode = mode; + } + + pub fn tool_call_dispatch_mode(&self) -> ToolCallDispatchMode { + self.tool_call_dispatch_mode + } + + /// Apply configuration. + pub fn config(mut self, config: EngineConfig) -> Self { + self.tool_call_dispatch_mode = config.tool_call_dispatch; self } @@ -2640,6 +3341,7 @@ impl Engine { llm_call_count: self.llm_call_count, tool_execution_batch_count: self.tool_execution_batch_count, max_turns: self.max_turns, + tool_call_dispatch_mode: self.tool_call_dispatch_mode, tool_execution_policy: self.tool_execution_policy, turn_start_cbs: self.turn_start_cbs, turn_end_cbs: self.turn_end_cbs, @@ -2823,6 +3525,7 @@ impl Engine { llm_call_count: self.llm_call_count, tool_execution_batch_count: self.tool_execution_batch_count, max_turns: self.max_turns, + tool_call_dispatch_mode: self.tool_call_dispatch_mode, tool_execution_policy: self.tool_execution_policy, turn_start_cbs: self.turn_start_cbs, turn_end_cbs: self.turn_end_cbs, diff --git a/crates/agen/src/lib.rs b/crates/agen/src/lib.rs index 89bdaee91..869d5c322 100644 --- a/crates/agen/src/lib.rs +++ b/crates/agen/src/lib.rs @@ -22,7 +22,7 @@ pub use agen_macros::{description, tool, tool_registry}; pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope}; pub use engine::{ Engine, EngineConfig, EngineError, EngineResult, EngineRunExit, EngineRunOutput, - LlmRetryNotice, RunInterruptionReason, ToolRegistryError, + LlmRetryNotice, RunInterruptionReason, ToolCallDispatchMode, ToolRegistryError, }; pub use handler::ToolUseBlockStart; pub use history::{History, HistoryEntry}; diff --git a/crates/agen/src/llm_client/client.rs b/crates/agen/src/llm_client/client.rs index d21acac85..b112d4291 100644 --- a/crates/agen/src/llm_client/client.rs +++ b/crates/agen/src/llm_client/client.rs @@ -38,6 +38,16 @@ impl std::fmt::Display for ConfigWarning { pub type ResponseStream = Pin> + Send>>; +/// Whether normalized `BlockStop(ToolUse)` events prove one call is complete. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum ToolCallCompletionSupport { + /// Individual tool-call completion events are provider-authoritative. + PerBlock, + /// Tool calls are complete only once the response stream itself completes. + #[default] + ResponseComplete, +} + /// LLMクライアントのtrait /// /// 各プロバイダはこのtraitを実装し、統一されたインターフェースを提供する。 @@ -59,6 +69,11 @@ pub trait LlmClient: Send + Sync { /// compaction) without access to the original construction parameters. fn clone_boxed(&self) -> Box; + /// Describe the strongest safe tool-call completion boundary. + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + ToolCallCompletionSupport::ResponseComplete + } + /// 設定をバリデーションし、未サポートの設定があれば警告を返す /// /// # Arguments @@ -92,6 +107,10 @@ impl LlmClient for Box { (**self).clone_boxed() } + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + (**self).tool_call_completion_support() + } + fn validate_config(&self, config: &RequestConfig) -> Vec { (**self).validate_config(config) } diff --git a/crates/agen/src/llm_client/scheme/anthropic/scheme_impl.rs b/crates/agen/src/llm_client/scheme/anthropic/scheme_impl.rs index ce15baa98..f4b7a0098 100644 --- a/crates/agen/src/llm_client/scheme/anthropic/scheme_impl.rs +++ b/crates/agen/src/llm_client/scheme/anthropic/scheme_impl.rs @@ -9,6 +9,7 @@ use crate::llm_client::{ ClientError, auth::AuthRequirement, capability::ModelCapability, + client::ToolCallCompletionSupport, event::{BlockType, Event, ReasoningBlockData}, scheme::Scheme, types::Request, @@ -104,4 +105,8 @@ impl Scheme for AnthropicScheme { fn default_capability(&self) -> ModelCapability { super::capability::default_capability() } + + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + ToolCallCompletionSupport::PerBlock + } } diff --git a/crates/agen/src/llm_client/scheme/mod.rs b/crates/agen/src/llm_client/scheme/mod.rs index 8306b32db..0ea78b819 100644 --- a/crates/agen/src/llm_client/scheme/mod.rs +++ b/crates/agen/src/llm_client/scheme/mod.rs @@ -16,7 +16,7 @@ use serde_json::Value; use super::auth::AuthRequirement; use super::capability::ModelCapability; -use super::client::ConfigWarning; +use super::client::{ConfigWarning, ToolCallCompletionSupport}; use super::error::ClientError; use super::event::Event; use super::types::{Request, RequestConfig}; @@ -81,6 +81,11 @@ pub trait Scheme: Clone + Send + Sync + 'static { /// 高レベルの client 構築層の責務で、scheme はここには関与しない。 fn default_capability(&self) -> ModelCapability; + /// Strongest safe completion boundary for streamed tool calls. + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + ToolCallCompletionSupport::ResponseComplete + } + /// scheme 側でサポートしていない `RequestConfig` フィールドを /// 警告として返す（例: OpenAI Chat は `top_k` 非対応）。 /// デフォルトは空 Vec。 @@ -89,3 +94,29 @@ pub trait Scheme: Clone + Send + Sync + 'static { Vec::new() } } + +#[cfg(test)] +mod completion_support_tests { + use super::*; + use crate::llm_client::ToolCallCompletionSupport; + + #[test] + fn only_schemes_with_authoritative_per_call_done_events_enable_early_dispatch() { + assert_eq!( + anthropic::AnthropicScheme::new().tool_call_completion_support(), + ToolCallCompletionSupport::PerBlock + ); + assert_eq!( + openai_responses::OpenAIResponsesScheme::new().tool_call_completion_support(), + ToolCallCompletionSupport::PerBlock + ); + assert_eq!( + openai_chat::OpenAIScheme::new().tool_call_completion_support(), + ToolCallCompletionSupport::ResponseComplete + ); + assert_eq!( + gemini::GeminiScheme::new().tool_call_completion_support(), + ToolCallCompletionSupport::ResponseComplete + ); + } +} diff --git a/crates/agen/src/llm_client/scheme/openai_responses/scheme_impl.rs b/crates/agen/src/llm_client/scheme/openai_responses/scheme_impl.rs index 8a6e4e0a2..0c3679c22 100644 --- a/crates/agen/src/llm_client/scheme/openai_responses/scheme_impl.rs +++ b/crates/agen/src/llm_client/scheme/openai_responses/scheme_impl.rs @@ -6,7 +6,7 @@ use crate::llm_client::{ ClientError, auth::AuthRequirement, capability::ModelCapability, - client::ConfigWarning, + client::{ConfigWarning, ToolCallCompletionSupport}, event::Event, scheme::Scheme, types::{Request, RequestConfig}, @@ -56,6 +56,10 @@ impl Scheme for OpenAIResponsesScheme { super::capability::default_capability() } + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + ToolCallCompletionSupport::PerBlock + } + fn validate_config(&self, config: &RequestConfig) -> Vec { let mut warnings = Vec::new(); // Some compatible backends reject `max_output_tokens` with HTTP 400. diff --git a/crates/agen/src/llm_client/transport.rs b/crates/agen/src/llm_client/transport.rs index 1198a9e2a..8eca250f1 100644 --- a/crates/agen/src/llm_client/transport.rs +++ b/crates/agen/src/llm_client/transport.rs @@ -510,6 +510,10 @@ impl LlmClient for HttpTransport { self.scheme.validate_config(config) } + fn tool_call_completion_support(&self) -> super::ToolCallCompletionSupport { + self.scheme.tool_call_completion_support() + } + async fn stream(&self, request: Request) -> Result { let total_started = Instant::now(); let path = self.scheme.path(&self.model_id); diff --git a/crates/agen/src/timeline/timeline.rs b/crates/agen/src/timeline/timeline.rs index f972ec31a..fcdac1fc3 100644 --- a/crates/agen/src/timeline/timeline.rs +++ b/crates/agen/src/timeline/timeline.rs @@ -256,8 +256,8 @@ where H: Handler, { handler: H, - scope: Option, - current_tool: Option<(String, String)>, // (id, name) + scopes: HashMap, + tools: HashMap, // index -> (id, name) } impl ToolUseBlockHandlerWrapper @@ -267,8 +267,8 @@ where fn new(handler: H) -> Self { Self { handler, - scope: None, - current_tool: None, + scopes: HashMap::new(), + tools: HashMap::new(), } } } @@ -279,60 +279,55 @@ where H::Scope: Send + Sync, { fn dispatch_start(&mut self, start: &BlockStart) { - if let Some(scope) = &mut self.scope { - if let BlockMetadata::ToolUse { id, name } = &start.metadata { - self.current_tool = Some((id.clone(), name.clone())); - self.handler.on_event( - scope, - &ToolUseBlockEvent::Start(ToolUseBlockStart { - index: start.index, - id: id.clone(), - name: name.clone(), - }), - ); - } + if let BlockMetadata::ToolUse { id, name } = &start.metadata { + let scope = self.scopes.entry(start.index).or_default(); + self.tools.insert(start.index, (id.clone(), name.clone())); + self.handler.on_event( + scope, + &ToolUseBlockEvent::Start(ToolUseBlockStart { + index: start.index, + id: id.clone(), + name: name.clone(), + }), + ); } } fn dispatch_delta(&mut self, delta: &BlockDelta) { - if let Some(scope) = &mut self.scope { - if let DeltaContent::InputJson(json) = &delta.delta { - self.handler - .on_event(scope, &ToolUseBlockEvent::InputJsonDelta(json.clone())); - } + if let DeltaContent::InputJson(json) = &delta.delta { + let scope = self.scopes.entry(delta.index).or_default(); + self.handler + .on_event(scope, &ToolUseBlockEvent::InputJsonDelta(json.clone())); } } fn dispatch_stop(&mut self, stop: &BlockStop) { - if let Some(scope) = &mut self.scope { - if let Some((id, name)) = self.current_tool.take() { - self.handler.on_event( - scope, - &ToolUseBlockEvent::Stop(ToolUseBlockStop { - index: stop.index, - id, - name, - }), - ); - } + if let (Some(mut scope), Some((id, name))) = ( + self.scopes.remove(&stop.index), + self.tools.remove(&stop.index), + ) { + self.handler.on_event( + &mut scope, + &ToolUseBlockEvent::Stop(ToolUseBlockStop { + index: stop.index, + id, + name, + }), + ); } } - fn dispatch_abort(&mut self, _abort: &BlockAbort) { - self.current_tool = None; + fn dispatch_abort(&mut self, abort: &BlockAbort) { + self.scopes.remove(&abort.index); + self.tools.remove(&abort.index); } - fn start_scope(&mut self) { - self.scope = Some(H::Scope::default()); - } + fn start_scope(&mut self) {} - fn end_scope(&mut self) { - self.scope = None; - self.current_tool = None; - } + fn end_scope(&mut self) {} fn has_scope(&self) -> bool { - self.scope.is_some() + !self.scopes.is_empty() } } diff --git a/crates/agen/src/timeline/tool_call_collector.rs b/crates/agen/src/timeline/tool_call_collector.rs index fa04b0ede..3c2ed0ca2 100644 --- a/crates/agen/src/timeline/tool_call_collector.rs +++ b/crates/agen/src/timeline/tool_call_collector.rs @@ -5,7 +5,6 @@ use crate::{ handler::{Handler, ToolUseBlockEvent, ToolUseBlockKind}, - llm_client::types::parse_tool_arguments, tool::ToolCall, }; use std::sync::{Arc, Mutex}; @@ -85,7 +84,18 @@ impl Handler for ToolCallCollector { // ブロック完了時にToolCallを確定 if let (Some(id), Some(name)) = (scope.current_id.take(), scope.current_name.take()) { - let input = parse_tool_arguments(&scope.input_json_buffer); + // A Stop event proves the provider finished the block, but + // malformed/non-object arguments are still not executable. + // Preserve an explicit null sentinel so the engine can commit + // a terminal validation error without guessing `{}`. + let input = if scope.input_json_buffer.trim().is_empty() { + serde_json::Value::Object(serde_json::Map::new()) + } else { + serde_json::from_str::(&scope.input_json_buffer) + .ok() + .filter(serde_json::Value::is_object) + .unwrap_or(serde_json::Value::Null) + }; let tool_call = ToolCall { id, name, input }; @@ -144,6 +154,43 @@ mod tests { ); } + #[test] + fn test_collect_interleaved_tool_call_indexes_and_ignore_duplicate_stop() { + let collector = ToolCallCollector::new(); + let mut timeline = Timeline::new(); + timeline.on_tool_use_block(collector.clone()); + + timeline.dispatch(&Event::tool_use_start(0, "call_0", "tool_0")); + timeline.dispatch(&Event::tool_input_delta(0, r#"{"a":"#)); + timeline.dispatch(&Event::tool_use_start(1, "call_1", "tool_1")); + timeline.dispatch(&Event::tool_input_delta(1, r#"{"b":2}"#)); + timeline.dispatch(&Event::tool_input_delta(0, r#"1}"#)); + timeline.dispatch(&Event::tool_use_stop(1)); + timeline.dispatch(&Event::tool_use_stop(1)); + timeline.dispatch(&Event::tool_use_stop(0)); + + let calls = collector.take_collected(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].id, "call_1"); + assert_eq!(calls[0].input, serde_json::json!({"b": 2})); + assert_eq!(calls[1].id, "call_0"); + assert_eq!(calls[1].input, serde_json::json!({"a": 1})); + } + + #[test] + fn malformed_completed_arguments_are_not_normalized_to_executable_object() { + let collector = ToolCallCollector::new(); + let mut timeline = Timeline::new(); + timeline.on_tool_use_block(collector.clone()); + timeline.dispatch(&Event::tool_use_start(0, "bad", "mutate")); + timeline.dispatch(&Event::tool_input_delta(0, r#"{"unterminated":"#)); + timeline.dispatch(&Event::tool_use_stop(0)); + + let calls = collector.take_collected(); + assert_eq!(calls.len(), 1); + assert!(calls[0].input.is_null()); + } + #[test] fn test_collect_multiple_tool_calls() { let collector = ToolCallCollector::new(); diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index 5df9f386e..f38f9f7cc 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -11,12 +11,16 @@ use agen::interceptor::{ PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo, }; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; +use agen::llm_client::{ + ClientError, LlmClient, Request, ResponseStream, ToolCallCompletionSupport, +}; use agen::tool::{ Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult, ToolResultDisposition, }; use agen::{ - Engine, EngineError, EngineRunExit, History, Item, RunInterruptionReason, ToolExecutionPolicy, + Engine, EngineError, EngineRunExit, History, Item, RunInterruptionReason, ToolCallDispatchMode, + ToolExecutionPolicy, }; use async_trait::async_trait; @@ -251,10 +255,476 @@ impl Tool for ContextRecordingTool { } } +#[derive(Clone)] +struct ControlledStreamClient { + receiver: Arc>>>>, + completion_support: ToolCallCompletionSupport, +} + +impl ControlledStreamClient { + fn new( + completion_support: ToolCallCompletionSupport, + ) -> ( + Self, + tokio::sync::mpsc::UnboundedSender>, + ) { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + ( + Self { + receiver: Arc::new(Mutex::new(Some(rx))), + completion_support, + }, + tx, + ) + } +} + +#[async_trait] +impl LlmClient for ControlledStreamClient { + async fn stream(&self, _request: Request) -> Result { + let mut rx = self + .receiver + .lock() + .unwrap() + .take() + .ok_or_else(|| ClientError::Api { + status: Some(500), + code: Some("controlled_stream_exhausted".to_string()), + message: "controlled stream already consumed".to_string(), + retry_after: None, + })?; + Ok(Box::pin(futures::stream::poll_fn(move |cx| { + rx.poll_recv(cx) + }))) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + self.completion_support + } +} + +#[derive(Clone)] +struct BarrierTool { + name: &'static str, + starts: Arc, + started: Arc, + release: Arc, +} + +impl BarrierTool { + fn new(name: &'static str) -> Self { + Self { + name, + starts: Arc::new(AtomicUsize::new(0)), + started: Arc::new(tokio::sync::Notify::new()), + release: Arc::new(tokio::sync::Notify::new()), + } + } + + fn definition(&self) -> ToolDefinition { + let tool = self.clone(); + Arc::new(move || { + ( + ToolMeta::new(tool.name).input_schema(serde_json::json!({"type":"object"})), + Arc::new(tool.clone()) as Arc, + ) + }) + } +} + +#[async_trait] +impl Tool for BarrierTool { + async fn execute( + &self, + _input_json: &str, + _ctx: ToolExecutionContext, + ) -> Result { + self.starts.fetch_add(1, Ordering::SeqCst); + self.started.notify_one(); + self.release.notified().await; + Ok(format!("{} complete", self.name).into()) + } +} + // ============================================================================= // Tests // ============================================================================= +#[tokio::test] +async fn default_dispatch_waits_for_response_completion() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = BarrierTool::new("barrier_default"); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_max_turns(Some(1)); + engine.register_tool(tool.definition()); + + let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); + engine.on_stream_event(move |_, _, event| { + if matches!(event, Event::Ping(_)) { + let _ = observed_tx.send(()); + } + }); + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + + tx.send(Ok(Event::tool_use_start( + 0, + "call_default", + "barrier_default", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tx.send(Ok(Event::ping())).unwrap(); + observed_rx.recv().await.unwrap(); + assert_eq!(probe.starts.load(Ordering::SeqCst), 0); + + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + tokio::time::timeout(Duration::from_secs(1), probe.started.notified()) + .await + .expect("tool starts after response completion"); + probe.release.notify_one(); + let _ = run.await.unwrap(); + assert_eq!(probe.starts.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_result() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let first = BarrierTool::new("barrier_a"); + let second = BarrierTool::new("barrier_b"); + let first_probe = first.clone(); + let second_probe = second.clone(); + let mut engine = Engine::new(client); + engine.set_max_turns(Some(1)); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(first.definition()); + engine.register_tool(second.definition()); + + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run both").await; + (output, history) + }); + + tx.send(Ok(Event::tool_use_start(0, "call_a", "barrier_a"))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tokio::time::timeout(Duration::from_secs(1), first_probe.started.notified()) + .await + .expect("first tool starts while stream is open"); + + tx.send(Ok(Event::tool_use_start(1, "call_b", "barrier_b"))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(1, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(1))).unwrap(); + tokio::time::timeout(Duration::from_secs(1), second_probe.started.notified()) + .await + .expect("second tool starts before first completes"); + assert_eq!(first_probe.starts.load(Ordering::SeqCst), 1); + assert_eq!(second_probe.starts.load(Ordering::SeqCst), 1); + + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + first_probe.release.notify_one(); + second_probe.release.notify_one(); + let (_output, history) = run.await.unwrap(); + + let calls = history + .iter() + .filter(|entry| matches!(entry.item, Item::ToolCall { .. })) + .count(); + let results = history + .iter() + .filter(|entry| matches!(entry.item, Item::ToolResult { .. })) + .count(); + assert_eq!((calls, results), (2, 2)); +} + +#[tokio::test] +async fn later_history_failure_terminalizes_already_started_early_sibling() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let first = BarrierTool::new("history_first"); + let second = BarrierTool::new("history_second"); + let first_probe = first.clone(); + let second_probe = second.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_tool_execution_policy(ToolExecutionPolicy { + pause_safe_boundary_timeout: Duration::from_millis(1), + cancellation_request_timeout: Duration::from_millis(1), + terminal_confirmation_timeout: Duration::from_millis(1), + }); + engine.register_tool(first.definition()); + engine.register_tool(second.definition()); + engine.on_history_append(|item| match item { + Item::ToolCall { call_id, .. } if call_id == "call_history_second" => { + Err("durable append rejected".to_string()) + } + _ => Ok(()), + }); + + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + tx.send(Ok(Event::tool_use_start( + 0, + "call_history_first", + "history_first", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tokio::time::timeout(Duration::from_secs(1), first_probe.started.notified()) + .await + .expect("first side effect starts"); + tx.send(Ok(Event::tool_use_start( + 1, + "call_history_second", + "history_second", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(1, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(1))).unwrap(); + drop(tx); + + let (output, history) = run.await.unwrap(); + assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + EngineError::HistoryAppend(ref message) + )) if message == "durable append rejected" + )); + assert_eq!(first_probe.starts.load(Ordering::SeqCst), 1); + assert_eq!(second_probe.starts.load(Ordering::SeqCst), 0); + assert!(history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call_history_first" + ))); +} + +#[tokio::test] +async fn malformed_completed_arguments_commit_error_without_executing_tool() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = SlowTool::new("must_not_run", 0); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_max_turns(Some(1)); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(tool.definition()); + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + + tx.send(Ok(Event::tool_use_start( + 0, + "call_malformed", + "must_not_run", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{"broken":"#))) + .unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + + let (_output, history) = run.await.unwrap(); + assert_eq!(probe.call_count(), 0); + assert!(history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { + call_id, + summary, + .. + } if call_id == "call_malformed" && summary.contains("complete JSON object") + ))); +} + +#[tokio::test] +async fn pause_during_early_execution_drains_result_and_resume_does_not_rerun() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = BarrierTool::new("barrier_pause"); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(tool.definition()); + let pause = engine.pause_sender(); + + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + tx.send(Ok(Event::tool_use_start(0, "call_pause", "barrier_pause"))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tokio::time::timeout(Duration::from_secs(1), probe.started.notified()) + .await + .expect("tool starts before pause"); + pause.send(()).await.unwrap(); + probe.release.notify_one(); + drop(tx); + + let (output, mut history) = run.await.unwrap(); + assert!(matches!(output.result, EngineRunExit::Paused)); + assert!(history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { call_id, .. } if call_id == "call_pause" + ))); + let mut engine = output.engine; + let _ = engine.resume(&mut history).await; + assert_eq!(probe.starts.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn interrupted_stream_after_early_start_terminalizes_without_continuation_or_reexecution() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = BarrierTool::new("barrier_interrupted"); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(tool.definition()); + + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + tx.send(Ok(Event::tool_use_start( + 0, + "call_interrupted", + "barrier_interrupted", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tokio::time::timeout(Duration::from_secs(1), probe.started.notified()) + .await + .expect("tool starts before stream interruption"); + tx.send(Err(ClientError::Api { + status: None, + code: Some("connection_lost".to_string()), + message: "connection lost".to_string(), + retry_after: None, + })) + .unwrap(); + probe.release.notify_one(); + drop(tx); + + let (output, mut history) = run.await.unwrap(); + assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(EngineError::Client( + ClientError::Api { + code: Some(ref code), + .. + } + ))) if code == "early_tool_stream_interrupted" + )); + assert_eq!(probe.starts.load(Ordering::SeqCst), 1); + assert_eq!( + history + .iter() + .filter(|entry| matches!(&entry.item, Item::ToolCall { call_id, .. } if call_id == "call_interrupted")) + .count(), + 1 + ); + assert_eq!( + history + .iter() + .filter(|entry| matches!(&entry.item, Item::ToolResult { call_id, .. } if call_id == "call_interrupted")) + .count(), + 1 + ); + + let mut engine = output.engine; + let _ = engine.resume(&mut history).await; + assert_eq!( + probe.starts.load(Ordering::SeqCst), + 1, + "terminalized early calls must not be resumed" + ); +} + +#[tokio::test] +async fn unsupported_provider_visibly_defers_opt_in_dispatch() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::ResponseComplete); + let tool = BarrierTool::new("barrier_fallback"); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_max_turns(Some(1)); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(tool.definition()); + let warnings = Arc::new(Mutex::new(Vec::new())); + let warning_probe = warnings.clone(); + engine.on_warning(move |warning| warning_probe.lock().unwrap().push(warning.to_string())); + let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); + engine.on_stream_event(move |_, _, event| { + if matches!(event, Event::Ping(_)) { + let _ = observed_tx.send(()); + } + }); + + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + tx.send(Ok(Event::tool_use_start( + 0, + "call_fallback", + "barrier_fallback", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tx.send(Ok(Event::ping())).unwrap(); + observed_rx.recv().await.unwrap(); + assert_eq!(probe.starts.load(Ordering::SeqCst), 0); + assert!(warnings.lock().unwrap().iter().any(|warning| { + warning.contains("provider does not expose a trustworthy per-call completion boundary") + })); + + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + tokio::time::timeout(Duration::from_secs(1), probe.started.notified()) + .await + .expect("fallback starts tool only after response completion"); + probe.release.notify_one(); + let _ = run.await.unwrap(); +} + /// Verify that multiple tools are executed in parallel /// /// If each tool takes 100ms, sequential execution would take 300ms+, diff --git a/crates/manifest/src/config.rs b/crates/manifest/src/config.rs index 9dc0e982e..d291883ef 100644 --- a/crates/manifest/src/config.rs +++ b/crates/manifest/src/config.rs @@ -542,6 +542,8 @@ pub struct EngineManifestConfig { #[serde(default)] pub reasoning: Option, #[serde(default)] + pub tool_call_dispatch: Option, + #[serde(default)] pub tool_output: ToolOutputLimitsPartial, #[serde(default)] pub file_upload: FileUploadLimitsPartial, @@ -868,6 +870,7 @@ impl EngineManifestConfig { top_k: upper.top_k.or(self.top_k), stop_sequences: upper.stop_sequences.or(self.stop_sequences), reasoning: upper.reasoning.or(self.reasoning), + tool_call_dispatch: upper.tool_call_dispatch.or(self.tool_call_dispatch), tool_output: self.tool_output.merge(upper.tool_output), file_upload: self.file_upload.merge(upper.file_upload), } @@ -1171,6 +1174,7 @@ impl TryFrom for WorkerManifest { top_k: cfg.engine.top_k, stop_sequences: cfg.engine.stop_sequences.unwrap_or_default(), reasoning: cfg.engine.reasoning, + tool_call_dispatch: cfg.engine.tool_call_dispatch.unwrap_or_default(), tool_output: ToolOutputLimits { default_max_bytes: cfg .engine @@ -1343,6 +1347,39 @@ mod tests { } } + #[test] + fn tool_call_dispatch_defaults_merges_and_round_trips() { + let manifest: WorkerManifest = minimal_valid().try_into().unwrap(); + assert_eq!( + manifest.engine.tool_call_dispatch, + crate::ToolCallDispatchMode::AfterResponse + ); + + let lower = WorkerManifestConfig { + engine: EngineManifestConfig { + tool_call_dispatch: Some(crate::ToolCallDispatchMode::OnToolCallComplete), + ..Default::default() + }, + ..Default::default() + }; + let upper = WorkerManifestConfig::default(); + assert_eq!( + lower.merge(upper).engine.tool_call_dispatch, + Some(crate::ToolCallDispatchMode::OnToolCallComplete) + ); + + let mut configured = minimal_valid(); + configured.engine.tool_call_dispatch = + Some(crate::ToolCallDispatchMode::OnToolCallComplete); + let manifest: WorkerManifest = configured.try_into().unwrap(); + let snapshot = serde_json::to_value(&manifest).unwrap(); + let restored: WorkerManifest = serde_json::from_value(snapshot).unwrap(); + assert_eq!( + restored.engine.tool_call_dispatch, + crate::ToolCallDispatchMode::OnToolCallComplete + ); + } + #[test] fn resolve_minimal_succeeds() { let manifest: WorkerManifest = minimal_valid().try_into().unwrap(); diff --git a/crates/manifest/src/lib.rs b/crates/manifest/src/lib.rs index e91c5b913..4cb678470 100644 --- a/crates/manifest/src/lib.rs +++ b/crates/manifest/src/lib.rs @@ -587,6 +587,20 @@ pub struct WorkerMeta { pub name: String, } +/// When the engine may begin executing provider-confirmed tool calls. +/// +/// `OnToolCallComplete` advances only the start time: calls may still overlap, +/// so completion order and externally allocated identifiers are not ordered. +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ToolCallDispatchMode { + /// Preserve the historical behavior: wait for the whole response stream. + #[default] + AfterResponse, + /// Start each call once the provider emits a trustworthy per-call completion. + OnToolCallComplete, +} + /// Worker-level configuration embedded in the manifest. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct EngineManifest { @@ -614,6 +628,9 @@ pub struct EngineManifest { pub stop_sequences: Vec, #[serde(default)] pub reasoning: Option, + /// Tool-call start timing. Omitted manifests preserve `after_response`. + #[serde(default)] + pub tool_call_dispatch: ToolCallDispatchMode, /// Byte-size caps applied to tool `content` before it reaches the /// conversation history. The section is optional in TOML — when /// omitted, `ToolOutputLimits::default()` (64 KiB default cap, no diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index f7bd11149..74a17ac47 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -965,6 +965,7 @@ fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfi stop_sequences: (!manifest.engine.stop_sequences.is_empty()) .then_some(manifest.engine.stop_sequences.clone()), reasoning: manifest.engine.reasoning.clone(), + tool_call_dispatch: Some(manifest.engine.tool_call_dispatch), tool_output: ToolOutputLimitsPartial { default_max_bytes: Some(manifest.engine.tool_output.default_max_bytes), per_tool: manifest.engine.tool_output.per_tool.clone(), @@ -2226,7 +2227,8 @@ max_tokens = 3333 Some("reviewer"), &[("reviewer", "reviewer.toml", REVIEWER_PROFILE)], ); - let parent = parent_manifest(&parent_root, Some(&parent_root.join("deny"))); + let mut parent = parent_manifest(&parent_root, Some(&parent_root.join("deny"))); + parent.engine.tool_call_dispatch = manifest::ToolCallDispatchMode::OnToolCallComplete; let scope = vec![abs_rule(&delegated, Permission::Read)]; let profile_config = child_config_from_profile( @@ -2248,6 +2250,10 @@ max_tokens = 3333 SpawnProfileSelector::Inherit, ); + assert_eq!( + inherit_config.engine.tool_call_dispatch, + Some(manifest::ToolCallDispatchMode::OnToolCallComplete) + ); for config in [profile_config, inherit_config] { assert_eq!(config.scope.allow, scope); assert!(config.scope.deny.is_empty()); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index ef8bd428c..6ab90b0d0 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -7138,6 +7138,12 @@ pub fn apply_worker_manifest( ) { worker.set_request_config(request_config_from_engine_manifest(wm)); worker.set_max_turns(wm.max_turns.map(|n| n.get())); + worker.set_tool_call_dispatch_mode(match wm.tool_call_dispatch { + manifest::ToolCallDispatchMode::AfterResponse => agen::ToolCallDispatchMode::AfterResponse, + manifest::ToolCallDispatchMode::OnToolCallComplete => { + agen::ToolCallDispatchMode::OnToolCallComplete + } + }); // Worker owns the lifecycle strategy for already-started tool operations. // The provider must first accept cooperative cancellation, then confirm a // terminal result before this bounded deadline; Agen handles only the @@ -8721,6 +8727,7 @@ mod build_summary_prompt_tests { top_k: Some(40), stop_sequences: vec!["\n\n".into(), "".into()], reasoning: None, + tool_call_dispatch: manifest::ToolCallDispatchMode::OnToolCallComplete, tool_output: manifest::ToolOutputLimits::default(), file_upload: manifest::FileUploadLimits::default(), }; @@ -8732,6 +8739,13 @@ mod build_summary_prompt_tests { assert_eq!(config.top_p, Some(0.9)); assert_eq!(config.top_k, Some(40)); assert_eq!(config.stop_sequences, vec!["\n\n", ""]); + + let mut engine = Engine::new(NoopClient); + apply_worker_manifest(&mut engine, &manifest); + assert_eq!( + engine.tool_call_dispatch_mode(), + agen::ToolCallDispatchMode::OnToolCallComplete + ); } #[test] diff --git a/docs/manifest.toml b/docs/manifest.toml index 92b671e6e..274a792e8 100644 --- a/docs/manifest.toml +++ b/docs/manifest.toml @@ -121,6 +121,14 @@ ref = "anthropic/claude-sonnet-4-6" # reasoning = "medium" # reasoning = -1 +# 任意。デフォルト: "after_response"。 +# "on_tool_call_complete" は provider が個別 ToolCall の完成を保証する場合だけ、 +# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider では +# 応答完了まで明示的に遅延する。これは直列実行モードではなく、先行 call の完了を +# 待たず後続 call も開始するため、完了順・DB 反映順・採番順は保証しない。 +# 早期開始後に stream が切断された場合は自動 continuation せず、安全に terminalize して停止する。 +# tool_call_dispatch = "on_tool_call_complete" + # 任意。tool 実行 content の byte 長キャップ。 # セクション省略時は default_max_bytes = 64 * 1024、per_tool 空。 # [worker.tool_output] From 8b951d53e6ced7a3616d0c090520c83c5b2c68f0 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 11:26:25 +0900 Subject: [PATCH 02/11] feat: add external workdir command permissions --- Cargo.lock | 1 + crates/server-api/src/lib.rs | 31 +- crates/workdir/Cargo.toml | 1 + crates/workdir/src/external.rs | 123 +++- crates/workdir/src/external_local_tests.rs | 75 ++- crates/workdir/src/lib.rs | 30 + crates/workdir/src/local.rs | 146 +++- crates/workspace-server/src/latest_schema.sql | 5 +- crates/workspace-server/src/server.rs | 627 +++++++++++++++--- crates/workspace-server/src/store.rs | 167 ++++- crates/yoi/src/main.rs | 118 +++- crates/yoi/src/workdir_share.rs | 259 ++++++-- docs/development/external-workdir-sharing.md | 42 ++ openapi/server-api.json | 40 +- .../src/lib/generated/companion-api.ts | 2 +- .../src/lib/generated/repository-api.ts | 2 +- .../src/lib/generated/runtime-api.ts | 2 +- web/workspace/src/lib/generated/ticket-api.ts | 2 +- .../src/lib/generated/workdir-api.ts | 8 +- .../src/lib/generated/worker-launch-api.ts | 2 +- 20 files changed, 1445 insertions(+), 238 deletions(-) create mode 100644 docs/development/external-workdir-sharing.md diff --git a/Cargo.lock b/Cargo.lock index 6b8e51672..8053c9987 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6181,6 +6181,7 @@ version = "0.1.0" dependencies = [ "async-trait", "fs-operation", + "libc", "manifest", "reqwest", "schemars", diff --git a/crates/server-api/src/lib.rs b/crates/server-api/src/lib.rs index 4790c3f4e..825616bf3 100644 --- a/crates/server-api/src/lib.rs +++ b/crates/server-api/src/lib.rs @@ -5241,13 +5241,28 @@ impl From for WorkingDirectorySumma } } -#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq, JsonSchema)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, JsonSchema)] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] -#[serde(rename_all = "snake_case")] -pub enum ExternalWorkdirPermission { - #[default] - ReadOnly, - ReadWrite, +#[serde(deny_unknown_fields)] +pub struct ExternalWorkdirPermissions { + /// Read files and use read-only discovery operations such as Glob and Grep. + pub read: bool, + /// Create or update files with Write and Edit. Existing-file mutation also + /// requires `read` so callers can satisfy the read-before-write contract. + pub write: bool, + /// Execute commands as the sharing CLI user on the provider host. This is + /// explicit authority and does not imply either filesystem category. + pub command: bool, +} + +impl Default for ExternalWorkdirPermissions { + fn default() -> Self { + Self { + read: true, + write: false, + command: false, + } + } } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonSchema)] @@ -5261,7 +5276,7 @@ pub struct ExternalWorkdirGrantCreateRequest { #[cfg_attr(feature = "typescript", ts(optional, type = "number | null"))] #[schemars(range(min = 0, max = 9_007_199_254_740_991_u64))] pub ttl_seconds: Option, - pub permission: ExternalWorkdirPermission, + pub permissions: ExternalWorkdirPermissions, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonSchema)] @@ -5274,7 +5289,7 @@ pub struct ExternalWorkdirGrantResponse { pub working_directory_id: String, pub provider_instance_id: String, pub display_name: String, - pub permission: ExternalWorkdirPermission, + pub permissions: ExternalWorkdirPermissions, /// `None` means the grant has no automatic expiry. It remains bounded by /// explicit revoke, provider disconnect, and provider shutdown cleanup. #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/crates/workdir/Cargo.toml b/crates/workdir/Cargo.toml index b5b7a99c8..0afb851e7 100644 --- a/crates/workdir/Cargo.toml +++ b/crates/workdir/Cargo.toml @@ -11,6 +11,7 @@ http-client = ["dep:reqwest"] [dependencies] async-trait.workspace = true fs-operation.workspace = true +libc.workspace = true manifest.workspace = true reqwest = { version = "0.13", default-features = false, features = ["json", "rustls"], optional = true } schemars.workspace = true diff --git a/crates/workdir/src/external.rs b/crates/workdir/src/external.rs index 22ff45d83..a1b12ec06 100644 --- a/crates/workdir/src/external.rs +++ b/crates/workdir/src/external.rs @@ -26,6 +26,17 @@ pub const MAX_EXTERNAL_RESULT_ITEMS: usize = 10_000; const MAX_EXTERNAL_GREP_CONTEXT: usize = 100; const MAX_EXTERNAL_SCOPE_RULES: usize = 64; pub const MAX_EXTERNAL_WRITE_BYTES: usize = 512 * 1024; +pub const MAX_EXTERNAL_COMMAND_BYTES: usize = 64 * 1024; +pub const MAX_EXTERNAL_COMMAND_OUTPUT_BYTES: usize = 512 * 1024; +pub const MAX_EXTERNAL_COMMAND_TIMEOUT_SECS: u64 = 600; +const MAX_EXTERNAL_COMMAND_CURSOR: usize = 1_000_000; + +fn valid_command_handle(handle: &crate::CommandHandle) -> bool { + !handle.0.is_empty() + && handle.0.len() <= MAX_PROTOCOL_ID_BYTES + && handle.0.trim() == handle.0 + && !handle.0.chars().any(char::is_control) +} /// Fail-closed wire version for the External Workdir provider protocol. #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] @@ -123,10 +134,9 @@ fn is_root_relative(path: &WorkdirPath) -> bool { }) } -/// Validate the filesystem-only subset shared by read-only and read-write -/// External Workdirs. Capability enforcement remains a separate grant/session -/// boundary; command operations never enter the provider protocol. -fn validate_filesystem_operation(operation: &WorkdirSessionOperation) -> Result<(), String> { +/// Validate the bounded operation subset available to External Workdirs. +/// Capability enforcement remains a separate grant/session boundary. +fn validate_external_operation(operation: &WorkdirSessionOperation) -> Result<(), String> { let path_is_valid = match operation { WorkdirSessionOperation::AuthorizeScope(request) => { is_root_relative(&request.path) @@ -182,17 +192,34 @@ fn validate_filesystem_operation(operation: &WorkdirSessionOperation) -> Result< .checked_add(request.new_string.len()) .is_some_and(|size| size <= MAX_EXTERNAL_WRITE_BYTES) } - WorkdirSessionOperation::CommandStart(_) - | WorkdirSessionOperation::CommandStatus(_) - | WorkdirSessionOperation::CommandOutput(_) - | WorkdirSessionOperation::CommandCancel(_) => { - return Err("command operations are not available to an External Workdir".to_string()); + WorkdirSessionOperation::CommandStart(request) => { + !request.command.is_empty() + && request.command.len() <= MAX_EXTERNAL_COMMAND_BYTES + && (1..=MAX_EXTERNAL_COMMAND_TIMEOUT_SECS).contains(&request.timeout_secs) + && (1..=MAX_EXTERNAL_COMMAND_OUTPUT_BYTES).contains(&request.output_limit) + && is_root_relative(&request.cwd) + && request.spill_dir.is_none() + && request.tool_call_id.as_ref().is_none_or(|tool_call_id| { + !tool_call_id.is_empty() + && tool_call_id.len() <= MAX_PROTOCOL_ID_BYTES + && !tool_call_id.chars().any(char::is_control) + }) + } + WorkdirSessionOperation::CommandStatus(handle) + | WorkdirSessionOperation::CommandCancel(handle) => valid_command_handle(handle), + WorkdirSessionOperation::CommandOutput(request) => { + valid_command_handle(&request.handle) + && request.cursor <= MAX_EXTERNAL_COMMAND_CURSOR + && (1..=MAX_EXTERNAL_COMMAND_OUTPUT_BYTES).contains(&request.limit) } }; if path_is_valid { Ok(()) } else { - Err("External Workdir operation paths must be root-relative".to_string()) + Err( + "External Workdir operation exceeds protocol bounds or uses an invalid path" + .to_string(), + ) } } @@ -215,7 +242,7 @@ impl TryFrom for ExternalWorkdirOperation { type Error = String; fn try_from(operation: WorkdirSessionOperation) -> Result { - validate_filesystem_operation(&operation)?; + validate_external_operation(&operation)?; Ok(Self(operation)) } } @@ -239,7 +266,7 @@ fn result_paths_fit<'a>(paths: impl IntoIterator) -> boo .is_some_and(|retained| retained <= fs_operation::MAX_RESULT_PATH_BYTES) } -fn validate_filesystem_result(result: &WorkdirSessionOperationResult) -> Result<(), String> { +fn validate_external_result(result: &WorkdirSessionOperationResult) -> Result<(), String> { let result_is_valid = match result { WorkdirSessionOperationResult::AuthorizeScope | WorkdirSessionOperationResult::ScopeRulesOverlap { .. } => true, @@ -273,14 +300,16 @@ fn validate_filesystem_result(result: &WorkdirSessionOperationResult) -> Result< result.bytes_written <= MAX_EXTERNAL_WRITE_BYTES && result.replacements <= MAX_EXTERNAL_RESULT_ITEMS } - WorkdirSessionOperationResult::CommandStart(_) - | WorkdirSessionOperationResult::CommandStatus(_) - | WorkdirSessionOperationResult::CommandOutput(_) - | WorkdirSessionOperationResult::CommandCancel => { - return Err( - "command operation results are not available to an External Workdir".to_string(), - ); + WorkdirSessionOperationResult::CommandStart(handle) => valid_command_handle(handle), + WorkdirSessionOperationResult::CommandStatus(_) => true, + WorkdirSessionOperationResult::CommandOutput(output) => { + output.content.len() <= MAX_EXTERNAL_COMMAND_OUTPUT_BYTES + && output + .next_cursor + .is_none_or(|cursor| cursor <= MAX_EXTERNAL_COMMAND_CURSOR) + && output.output_path.is_none() } + WorkdirSessionOperationResult::CommandCancel => true, }; if result_is_valid { Ok(()) @@ -308,7 +337,7 @@ impl TryFrom for ExternalWorkdirOperationResult { type Error = String; fn try_from(result: WorkdirSessionOperationResult) -> Result { - validate_filesystem_result(&result)?; + validate_external_result(&result)?; Ok(Self(result)) } } @@ -544,7 +573,7 @@ mod tests { } #[test] - fn filesystem_protocol_rejects_host_paths_and_commands_but_allows_bounded_mutations() { + fn protocol_rejects_host_paths_and_unbounded_commands_but_allows_bounded_operations() { let absolute = WorkdirSessionOperation::Read(ReadRequest { path: WorkdirPath::new_scoped("/home/operator/private/session.log").unwrap(), offset: 0, @@ -563,6 +592,29 @@ mod tests { }); assert!(ExternalWorkdirOperation::try_from(command).is_err()); + let command = WorkdirSessionOperation::CommandStart(CommandRequest { + command: "printf bounded".to_string(), + timeout_secs: 30, + output_limit: 12 * 1024, + cwd: WorkdirPath::root(), + spill_dir: None, + tool_call_id: Some("tool-call-1".to_string()), + }); + assert!(ExternalWorkdirOperation::try_from(command).is_ok()); + assert!( + ExternalWorkdirOperation::try_from(WorkdirSessionOperation::CommandStart( + CommandRequest { + command: "x".repeat(MAX_EXTERNAL_COMMAND_BYTES + 1), + timeout_secs: 30, + output_limit: 12 * 1024, + cwd: WorkdirPath::root(), + spill_dir: None, + tool_call_id: None, + } + )) + .is_err() + ); + assert!( ExternalWorkdirOperation::try_from(WorkdirSessionOperation::Write(WriteRequest { path: WorkdirPath::new("created.txt").unwrap(), @@ -589,6 +641,35 @@ mod tests { })) .is_err() ); + + assert!( + ExternalWorkdirOperationResult::try_from(WorkdirSessionOperationResult::CommandOutput( + crate::CommandOutput { + status: crate::CommandStatus::Completed, + exit_code: Some(0), + timed_out: false, + content: "bounded".to_string(), + next_cursor: None, + truncated: false, + output_path: None, + } + )) + .is_ok() + ); + assert!( + ExternalWorkdirOperationResult::try_from(WorkdirSessionOperationResult::CommandOutput( + crate::CommandOutput { + status: crate::CommandStatus::Completed, + exit_code: Some(0), + timed_out: false, + content: "bounded".to_string(), + next_cursor: None, + truncated: true, + output_path: Some("/provider/private/output".into()), + } + )) + .is_err() + ); } #[test] diff --git a/crates/workdir/src/external_local_tests.rs b/crates/workdir/src/external_local_tests.rs index 943eb3dd5..664f8eca3 100644 --- a/crates/workdir/src/external_local_tests.rs +++ b/crates/workdir/src/external_local_tests.rs @@ -2,10 +2,10 @@ use crate::http::{ WorkdirSessionOperation, WorkdirSessionOperationResult, dispatch_workdir_session_operation, }; use crate::{ - BoundedReadLimits, CommandRequest, EditRequest, ExternalWorkdirRoot, GlobRequest, - GrepOutputMode, GrepRequest, LocalWorkdirSession, ReadRequest, Workdir, WorkdirError, - WorkdirPath, WorkdirSession, WorkdirSessionCapabilities, WorkdirSessionCapability, - WriteRequest, + BoundedReadLimits, CommandOutputRequest, CommandRequest, EditRequest, ExternalWorkdirRoot, + GlobRequest, GrepOutputMode, GrepRequest, LocalWorkdirSession, ReadRequest, Workdir, + WorkdirError, WorkdirPath, WorkdirSession, WorkdirSessionCapabilities, + WorkdirSessionCapability, WriteRequest, }; fn external_session(root: &tempfile::TempDir, limits: BoundedReadLimits) -> LocalWorkdirSession { @@ -162,6 +162,73 @@ async fn external_local_provider_read_write_is_filesystem_only() { assert!(!root.path().join("command-ran").exists()); } +#[tokio::test] +async fn external_command_permission_is_explicit_and_independent_from_file_access() { + let root = tempfile::tempdir().unwrap(); + std::fs::write(root.path().join("item.txt"), "private").unwrap(); + let pinned = ExternalWorkdirRoot::pin(root.path()).unwrap(); + let session = LocalWorkdirSession::external_with_capabilities_pinned( + Workdir::new("external-workdir-command"), + pinned, + BoundedReadLimits::new(4096, 1024).unwrap(), + WorkdirSessionCapabilities::COMMAND_ONLY, + ) + .unwrap(); + + assert_eq!( + session.capabilities(), + WorkdirSessionCapabilities::COMMAND_ONLY + ); + assert!(matches!( + session + .read(ReadRequest { + path: WorkdirPath::new("item.txt").unwrap(), + offset: 0, + limit: 10, + max_bytes: 1024, + }) + .await, + Err(WorkdirError::Unsupported(WorkdirSessionCapability::Read)) + )); + + let handle = session + .start_command(CommandRequest { + command: "printf command-ok".to_string(), + timeout_secs: 5, + output_limit: 1024, + cwd: WorkdirPath::root(), + spill_dir: None, + tool_call_id: None, + }) + .await + .unwrap(); + let output = session + .command_output(CommandOutputRequest { + handle, + cursor: 0, + limit: 1024, + wait: true, + }) + .await + .unwrap(); + assert_eq!(output.content, "command-ok"); + assert_eq!(output.exit_code, Some(0)); +} + +#[test] +fn external_write_category_never_implicitly_enables_read() { + let root = tempfile::tempdir().unwrap(); + let pinned = ExternalWorkdirRoot::pin(root.path()).unwrap(); + let error = LocalWorkdirSession::external_with_capabilities_pinned( + Workdir::new("external-workdir-invalid-write"), + pinned, + BoundedReadLimits::new(4096, 1024).unwrap(), + WorkdirSessionCapabilities::WRITE_ONLY, + ) + .unwrap_err(); + assert!(error.to_string().contains("WRITE requires READ")); +} + #[tokio::test] async fn common_dispatcher_preserves_operation_result_pairing_and_capabilities() { let root = tempfile::tempdir().unwrap(); diff --git a/crates/workdir/src/lib.rs b/crates/workdir/src/lib.rs index 70a0b5a62..806a1431a 100644 --- a/crates/workdir/src/lib.rs +++ b/crates/workdir/src/lib.rs @@ -134,6 +134,16 @@ impl WorkdirSessionCapabilities { bits: Self::READ | Self::GLOB | Self::GREP, }; + /// Operator-facing WRITE category without implicitly granting READ or COMMAND. + pub const WRITE_ONLY: Self = Self { + bits: Self::WRITE | Self::EDIT, + }; + + /// Operator-facing COMMAND category without implicitly granting file access. + pub const COMMAND_ONLY: Self = Self { + bits: Self::COMMAND, + }; + /// Filesystem read/write authority without process execution. /// /// External Workdirs use this capability set so operator-approved file @@ -391,3 +401,23 @@ impl From for WorkdirError { } } } + +#[cfg(test)] +mod capability_tests { + use super::*; + + fn bits(capabilities: WorkdirSessionCapabilities) -> u8 { + serde_json::to_value(capabilities).unwrap()["bits"] + .as_u64() + .unwrap() as u8 + } + + #[test] + fn external_categories_preserve_operation_bit_values() { + assert_eq!(bits(WorkdirSessionCapabilities::READ_ONLY), 25); + assert_eq!(bits(WorkdirSessionCapabilities::WRITE_ONLY), 6); + assert_eq!(bits(WorkdirSessionCapabilities::READ_WRITE), 31); + assert_eq!(bits(WorkdirSessionCapabilities::COMMAND_ONLY), 32); + assert_eq!(bits(WorkdirSessionCapabilities::ALL), 63); + } +} diff --git a/crates/workdir/src/local.rs b/crates/workdir/src/local.rs index 2d67a3a12..4a4a3b61e 100644 --- a/crates/workdir/src/local.rs +++ b/crates/workdir/src/local.rs @@ -11,6 +11,8 @@ use std::collections::{BTreeMap, HashMap}; use std::fmt::Debug; use std::io::{Read as _, Seek as _, SeekFrom, Write as _}; +#[cfg(unix)] +use std::os::unix::process::CommandExt as _; use std::path::{Path, PathBuf}; use std::process::Stdio; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; @@ -762,6 +764,42 @@ impl LocalWorkdirSession { Self::external_read_write_pinned(workdir, root, read_limits) } + /// Construct an External provider session with an exact operator-approved + /// READ/WRITE/COMMAND category set. WRITE requires READ so callers can + /// satisfy the read-before-write contract; COMMAND remains independent. + pub fn external_with_capabilities_pinned( + workdir: Workdir, + pinned: ExternalWorkdirRoot, + read_limits: BoundedReadLimits, + capabilities: WorkdirSessionCapabilities, + ) -> Result { + let read = WorkdirSessionCapabilities::READ_ONLY; + let read_write = WorkdirSessionCapabilities::READ_WRITE; + let command = WorkdirSessionCapabilities::COMMAND_ONLY; + if ![ + read, + read_write, + command, + read.union(command), + read_write.union(command), + ] + .contains(&capabilities) + { + return Err(WorkdirError::InvalidArgument( + "External Workdir capabilities must contain at least one of READ or COMMAND, and WRITE requires READ" + .to_string(), + )); + } + let permission = if capabilities.supports(WorkdirSessionCapability::Write) { + Permission::Write + } else { + // Command execution still needs provider-internal read access to + // validate and enter its cwd. This scope is not an exposed READ grant. + Permission::Read + }; + Self::external_pinned(workdir, pinned, read_limits, permission, capabilities) + } + /// Construct a read-only provider session from a root pinned before remote /// grant creation. pub fn external_read_only_pinned( @@ -769,11 +807,10 @@ impl LocalWorkdirSession { pinned: ExternalWorkdirRoot, read_limits: BoundedReadLimits, ) -> Result { - Self::external_pinned( + Self::external_with_capabilities_pinned( workdir, pinned, read_limits, - Permission::Read, WorkdirSessionCapabilities::READ_ONLY, ) } @@ -785,11 +822,10 @@ impl LocalWorkdirSession { pinned: ExternalWorkdirRoot, read_limits: BoundedReadLimits, ) -> Result { - Self::external_pinned( + Self::external_with_capabilities_pinned( workdir, pinned, read_limits, - Permission::Write, WorkdirSessionCapabilities::READ_WRITE, ) } @@ -1556,6 +1592,45 @@ fn sanitize_error(error: WorkdirError, logical: &WorkdirPath) -> WorkdirError { } } +#[derive(Debug)] +struct CommandProcessGroupGuard { + process_group_id: Option, +} + +impl CommandProcessGroupGuard { + fn new(process_group_id: Option) -> Self { + Self { process_group_id } + } + + fn terminate(&mut self) { + let Some(process_group_id) = self.process_group_id.take() else { + return; + }; + terminate_process_group(process_group_id); + } +} + +impl Drop for CommandProcessGroupGuard { + fn drop(&mut self) { + self.terminate(); + } +} + +#[cfg(unix)] +fn terminate_process_group(process_group_id: u32) { + let Ok(process_group_id) = i32::try_from(process_group_id) else { + return; + }; + // The command is spawned as its own process-group leader below. A negative + // pid addresses that exact group, including descendants that outlive bash. + unsafe { + libc::kill(-process_group_id, libc::SIGKILL); + } +} + +#[cfg(not(unix))] +fn terminate_process_group(_process_group_id: u32) {} + async fn run_command( cwd: PathBuf, request: CommandRequest, @@ -1574,7 +1649,8 @@ async fn run_command( .map_err(|error| WorkdirError::io(&stderr_path, error))?; telemetry.started(&command_id, request.tool_call_id.clone()); - let mut child = match Command::new("bash") + let mut command = Command::new("bash"); + command .arg("-c") .arg(&request.command) .current_dir(&cwd) @@ -1582,15 +1658,17 @@ async fn run_command( .stdin(Stdio::null()) .stdout(Stdio::from(stdout_file)) .stderr(Stdio::from(stderr_file)) - .kill_on_drop(true) - .spawn() - { + .kill_on_drop(true); + #[cfg(unix)] + command.as_std_mut().process_group(0); + let mut child = match command.spawn() { Ok(child) => child, Err(error) => { telemetry.terminal(&command_id, CommandStatus::Failed, None); return Err(WorkdirError::io(&cwd, error)); } }; + let mut process_group = CommandProcessGroupGuard::new(child.id()); let mut stdout_reader = std::fs::File::open(&stdout_path).map_err(|error| WorkdirError::io(&stdout_path, error))?; @@ -1609,18 +1687,21 @@ async fn run_command( tokio::select! { exit = child.wait() => { let exit = exit.map_err(|error| WorkdirError::io(&cwd, error))?; + process_group.terminate(); break ( if exit.success() { CommandStatus::Completed } else { CommandStatus::Failed }, exit.code(), ); } _ = &mut timeout => { + process_group.terminate(); let _ = child.start_kill(); let exit_code = child.wait().await.ok().and_then(|status| status.code()); break (CommandStatus::TimedOut, exit_code); } changed = cancel.changed() => { if changed.is_err() || *cancel.borrow() { + process_group.terminate(); let _ = child.start_kill(); let exit_code = child.wait().await.ok().and_then(|status| status.code()); break (CommandStatus::Cancelled, exit_code); @@ -3473,6 +3554,55 @@ mod tests { assert!(released.load(Ordering::Acquire)); } + #[cfg(unix)] + #[tokio::test] + async fn cancellation_kills_command_process_group_descendants() { + let dir = TempDir::new().unwrap(); + let workdir = make_fs(&dir); + let handle = WorkdirSession::start_command( + &workdir, + CommandRequest { + command: "(sleep 0.4; echo leaked > child-marker) & echo ready > child-ready; wait" + .into(), + timeout_secs: 60, + output_limit: 1024, + cwd: WorkdirPath::root(), + spill_dir: None, + tool_call_id: None, + }, + ) + .await + .unwrap(); + for _ in 0..100 { + if dir.path().join("child-ready").exists() { + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + assert!(dir.path().join("child-ready").exists()); + + WorkdirSession::cancel_command(&workdir, handle.clone()) + .await + .unwrap(); + let output = WorkdirSession::command_output( + &workdir, + CommandOutputRequest { + handle, + cursor: 0, + limit: 1024, + wait: true, + }, + ) + .await + .unwrap(); + assert_eq!(output.status, CommandStatus::Cancelled); + tokio::time::sleep(Duration::from_millis(500)).await; + assert!( + !dir.path().join("child-marker").exists(), + "cancelled command descendant survived its process group" + ); + } + #[tokio::test] async fn provider_cancels_active_command() { let dir = TempDir::new().unwrap(); diff --git a/crates/workspace-server/src/latest_schema.sql b/crates/workspace-server/src/latest_schema.sql index 674bee27d..0accd1d9a 100644 --- a/crates/workspace-server/src/latest_schema.sql +++ b/crates/workspace-server/src/latest_schema.sql @@ -697,7 +697,10 @@ CREATE TABLE external_workdir_grants ( workdir_id TEXT NOT NULL, provider_instance_id TEXT NOT NULL, display_name TEXT NOT NULL, - permissions TEXT NOT NULL CHECK (permissions IN ('read_only', 'read_write')), + permissions TEXT NOT NULL CHECK (permissions IN ( + 'read_only', 'read_write', 'command_only', + 'read_command', 'read_write_command' + )), created_by TEXT NOT NULL, created_at TEXT NOT NULL, expires_at TEXT, diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 69cfef1d0..57f68cb53 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -711,14 +711,40 @@ impl ExternalOperationCancelGuard { fn disarm(&mut self) { self.operation_id = None; } + + fn cancel(&mut self) { + let Some(operation_id) = self.operation_id.take() else { + return; + }; + enqueue_external_operation_cancel(self.sender.clone(), operation_id); + } } impl Drop for ExternalOperationCancelGuard { fn drop(&mut self) { - if let Some(operation_id) = self.operation_id.take() { - let _ = self - .sender - .try_send(ExternalProviderCommand::Cancel { operation_id }); + self.cancel(); + } +} + +fn enqueue_external_operation_cancel( + sender: tokio::sync::mpsc::Sender, + operation_id: ExternalWorkdirOperationId, +) { + let command = ExternalProviderCommand::Cancel { operation_id }; + match sender.try_send(command) { + Ok(()) | Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {} + Err(tokio::sync::mpsc::error::TrySendError::Full(command)) => { + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return; + }; + runtime.spawn(async move { + // Never delay the failed operation's return on a saturated provider + // queue, but give the provider a bounded opportunity to accept the + // cancellation as it drains already queued messages. + let _ = + tokio::time::timeout(std::time::Duration::from_secs(5), sender.send(command)) + .await; + }); } } } @@ -792,6 +818,8 @@ struct ExternalProviderWorkdirSession { connection: Arc, audit: Option, closed: std::sync::atomic::AtomicBool, + close_lock: tokio::sync::Mutex<()>, + active_commands: tokio::sync::Mutex>, } impl std::fmt::Debug for ExternalProviderWorkdirSession { @@ -815,6 +843,8 @@ impl ExternalProviderWorkdirSession { connection, audit, closed: std::sync::atomic::AtomicBool::new(false), + close_lock: tokio::sync::Mutex::new(()), + active_commands: tokio::sync::Mutex::new(HashMap::new()), } } @@ -858,6 +888,7 @@ impl ExternalProviderWorkdirSession { operation: WorkdirSessionOperation, ) -> std::result::Result { let operation_kind = external_workdir_operation_kind(&operation); + let operation_timeout = external_workdir_operation_timeout(&operation); if self.closed.load(Ordering::Acquire) { return Err(workdir::WorkdirError::SessionClosed); } @@ -914,7 +945,7 @@ impl ExternalProviderWorkdirSession { "External Workdir provider disconnected".to_string(), )); } - let result = match tokio::time::timeout(std::time::Duration::from_secs(30), receive).await { + let result = match tokio::time::timeout(operation_timeout, receive).await { Ok(result) => { cancel_guard.disarm(); match result { @@ -925,14 +956,7 @@ impl ExternalProviderWorkdirSession { } } Err(_) => { - if let Some(cancel_operation_id) = cancel_guard.operation_id.take() { - let _ = cancel_guard - .sender - .send(ExternalProviderCommand::Cancel { - operation_id: cancel_operation_id, - }) - .await; - } + cancel_guard.cancel(); Err(workdir::WorkdirError::Unavailable( "External Workdir provider operation timed out".to_string(), )) @@ -950,6 +974,17 @@ impl ExternalProviderWorkdirSession { result } + fn ensure_capability( + &self, + capability: workdir::WorkdirSessionCapability, + ) -> std::result::Result<(), workdir::WorkdirError> { + if self.connection.capabilities.supports(capability) { + Ok(()) + } else { + Err(workdir::WorkdirError::Unsupported(capability)) + } + } + fn mismatch(expected: &str) -> workdir::WorkdirError { workdir::WorkdirError::Unavailable(format!( "External Workdir provider returned a mismatched result; expected {expected}" @@ -957,6 +992,17 @@ impl ExternalProviderWorkdirSession { } } +fn external_workdir_operation_timeout(operation: &WorkdirSessionOperation) -> std::time::Duration { + if matches!( + operation, + WorkdirSessionOperation::CommandOutput(request) if request.wait + ) { + std::time::Duration::from_secs(workdir::external::MAX_EXTERNAL_COMMAND_TIMEOUT_SECS + 10) + } else { + std::time::Duration::from_secs(30) + } +} + fn external_workdir_operation_kind(operation: &WorkdirSessionOperation) -> &'static str { match operation { WorkdirSessionOperation::AuthorizeScope(_) => "authorize_scope", @@ -988,6 +1034,7 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { &self, request: workdir::WorkdirScopeAuthorizationRequest, ) -> std::result::Result<(), workdir::WorkdirError> { + self.ensure_capability(workdir::WorkdirSessionCapability::Read)?; match self .operate(WorkdirSessionOperation::AuthorizeScope(request)) .await? @@ -1000,6 +1047,7 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { &self, request: workdir::WorkdirScopeOverlapRequest, ) -> std::result::Result { + self.ensure_capability(workdir::WorkdirSessionCapability::Read)?; match self .operate(WorkdirSessionOperation::ScopeRulesOverlap(request)) .await? @@ -1012,6 +1060,7 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { &self, request: workdir::StatRequest, ) -> std::result::Result { + self.ensure_capability(workdir::WorkdirSessionCapability::Read)?; match self.operate(WorkdirSessionOperation::Stat(request)).await? { WorkdirSessionOperationResult::Stat(value) => Ok(value), _ => Err(Self::mismatch("stat")), @@ -1021,6 +1070,7 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { &self, request: workdir::ReadRequest, ) -> std::result::Result { + self.ensure_capability(workdir::WorkdirSessionCapability::Read)?; match self.operate(WorkdirSessionOperation::Read(request)).await? { WorkdirSessionOperationResult::Read(value) => Ok(value), _ => Err(Self::mismatch("read")), @@ -1030,15 +1080,7 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { &self, request: workdir::WriteRequest, ) -> std::result::Result { - if !self - .connection - .capabilities - .supports(workdir::WorkdirSessionCapability::Write) - { - return Err(workdir::WorkdirError::Unsupported( - workdir::WorkdirSessionCapability::Write, - )); - } + self.ensure_capability(workdir::WorkdirSessionCapability::Write)?; match self .operate(WorkdirSessionOperation::Write(request)) .await? @@ -1051,15 +1093,7 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { &self, request: workdir::EditRequest, ) -> std::result::Result { - if !self - .connection - .capabilities - .supports(workdir::WorkdirSessionCapability::Edit) - { - return Err(workdir::WorkdirError::Unsupported( - workdir::WorkdirSessionCapability::Edit, - )); - } + self.ensure_capability(workdir::WorkdirSessionCapability::Edit)?; match self.operate(WorkdirSessionOperation::Edit(request)).await? { WorkdirSessionOperationResult::Edit(value) => Ok(value), _ => Err(Self::mismatch("edit")), @@ -1069,6 +1103,7 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { &self, request: workdir::ListRequest, ) -> std::result::Result { + self.ensure_capability(workdir::WorkdirSessionCapability::Read)?; match self.operate(WorkdirSessionOperation::List(request)).await? { WorkdirSessionOperationResult::List(value) => Ok(value), _ => Err(Self::mismatch("list")), @@ -1078,6 +1113,7 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { &self, request: workdir::GlobRequest, ) -> std::result::Result { + self.ensure_capability(workdir::WorkdirSessionCapability::Glob)?; match self.operate(WorkdirSessionOperation::Glob(request)).await? { WorkdirSessionOperationResult::Glob(value) => Ok(value), _ => Err(Self::mismatch("glob")), @@ -1087,6 +1123,7 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { &self, request: workdir::GrepRequest, ) -> std::result::Result { + self.ensure_capability(workdir::WorkdirSessionCapability::Grep)?; match self.operate(WorkdirSessionOperation::Grep(request)).await? { WorkdirSessionOperationResult::Grep(value) => Ok(value), _ => Err(Self::mismatch("grep")), @@ -1094,37 +1131,107 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { } async fn start_command( &self, - _request: workdir::CommandRequest, + mut request: workdir::CommandRequest, ) -> std::result::Result { - Err(workdir::WorkdirError::Unsupported( - workdir::WorkdirSessionCapability::Command, - )) + 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; + match self + .operate(WorkdirSessionOperation::CommandStart(request)) + .await? + { + WorkdirSessionOperationResult::CommandStart(handle) => { + self.active_commands + .lock() + .await + .insert(handle.0.clone(), handle.clone()); + Ok(handle) + } + _ => Err(Self::mismatch("command_start")), + } } async fn command_status( &self, - _handle: workdir::CommandHandle, + handle: workdir::CommandHandle, ) -> std::result::Result { - Err(workdir::WorkdirError::Unsupported( - workdir::WorkdirSessionCapability::Command, - )) + self.ensure_capability(workdir::WorkdirSessionCapability::Command)?; + match self + .operate(WorkdirSessionOperation::CommandStatus(handle)) + .await? + { + WorkdirSessionOperationResult::CommandStatus(status) => Ok(status), + _ => Err(Self::mismatch("command_status")), + } } async fn command_output( &self, - _request: workdir::CommandOutputRequest, + request: workdir::CommandOutputRequest, ) -> std::result::Result { - Err(workdir::WorkdirError::Unsupported( - workdir::WorkdirSessionCapability::Command, - )) + self.ensure_capability(workdir::WorkdirSessionCapability::Command)?; + let command_id = request.handle.0.clone(); + match self + .operate(WorkdirSessionOperation::CommandOutput(request)) + .await? + { + WorkdirSessionOperationResult::CommandOutput(output) => { + if output.status != workdir::CommandStatus::Running && output.next_cursor.is_none() + { + self.active_commands.lock().await.remove(&command_id); + } + Ok(output) + } + _ => Err(Self::mismatch("command_output")), + } } async fn cancel_command( &self, - _handle: workdir::CommandHandle, + handle: workdir::CommandHandle, ) -> std::result::Result<(), workdir::WorkdirError> { - Err(workdir::WorkdirError::Unsupported( - workdir::WorkdirSessionCapability::Command, - )) + self.ensure_capability(workdir::WorkdirSessionCapability::Command)?; + match self + .operate(WorkdirSessionOperation::CommandCancel(handle)) + .await? + { + WorkdirSessionOperationResult::CommandCancel => Ok(()), + _ => Err(Self::mismatch("command_cancel")), + } } async fn close(&self) -> std::result::Result<(), workdir::WorkdirError> { + let _close_guard = self.close_lock.lock().await; + if self.closed.load(Ordering::Acquire) { + return Ok(()); + } + let commands = self + .active_commands + .lock() + .await + .drain() + .map(|(_, handle)| handle) + .collect::>(); + for handle in commands { + // Detach/session replacement is terminal authority for commands + // started through this exact provider generation. Request cancel, + // then wait for provider-owned process reaping before closing. + let _ = self + .operate(WorkdirSessionOperation::CommandCancel(handle.clone())) + .await; + let _ = self + .operate(WorkdirSessionOperation::CommandOutput( + workdir::CommandOutputRequest { + handle, + cursor: 0, + limit: workdir::external::MAX_EXTERNAL_COMMAND_OUTPUT_BYTES, + wait: true, + }, + )) + .await; + } self.closed.store(true, Ordering::Release); Ok(()) } @@ -3674,16 +3781,75 @@ struct ValidatedTicketImplementationTarget { capabilities: workdir::WorkdirSessionCapabilities, } -fn external_workdir_capabilities(permissions: &str) -> Result { - match permissions { - "read_only" => Ok(workdir::WorkdirSessionCapabilities::READ_ONLY), - "read_write" => Ok(workdir::WorkdirSessionCapabilities::READ_WRITE), - _ => Err(Error::Store( - "External Workdir grant has invalid persisted permissions".to_string(), +fn external_workdir_permissions(persisted: &str) -> Result { + let permissions = match persisted { + "read_only" => server_api::ExternalWorkdirPermissions { + read: true, + write: false, + command: false, + }, + "read_write" => server_api::ExternalWorkdirPermissions { + read: true, + write: true, + command: false, + }, + "command_only" => server_api::ExternalWorkdirPermissions { + read: false, + write: false, + command: true, + }, + "read_command" => server_api::ExternalWorkdirPermissions { + read: true, + write: false, + command: true, + }, + "read_write_command" => server_api::ExternalWorkdirPermissions { + read: true, + write: true, + command: true, + }, + _ => { + return Err(Error::Store( + "External Workdir grant has invalid persisted permissions".to_string(), + )); + } + }; + Ok(permissions) +} + +fn persist_external_workdir_permissions( + permissions: server_api::ExternalWorkdirPermissions, +) -> Result<&'static str> { + match (permissions.read, permissions.write, permissions.command) { + (true, false, false) => Ok("read_only"), + (true, true, false) => Ok("read_write"), + (false, false, true) => Ok("command_only"), + (true, false, true) => Ok("read_command"), + (true, true, true) => Ok("read_write_command"), + (false, true, _) => Err(Error::InvalidInput( + "External Workdir WRITE permission requires READ for read-before-write".to_string(), + )), + (false, false, false) => Err(Error::InvalidInput( + "External Workdir sharing requires at least one of READ or COMMAND".to_string(), )), } } +fn external_workdir_capabilities(permissions: &str) -> Result { + let permissions = external_workdir_permissions(permissions)?; + let mut capabilities = workdir::WorkdirSessionCapabilities::EMPTY; + if permissions.read { + capabilities = capabilities.union(workdir::WorkdirSessionCapabilities::READ_ONLY); + } + if permissions.write { + capabilities = capabilities.union(workdir::WorkdirSessionCapabilities::WRITE_ONLY); + } + if permissions.command { + capabilities = capabilities.union(workdir::WorkdirSessionCapabilities::COMMAND_ONLY); + } + Ok(capabilities) +} + fn workdir_source_capabilities( api: &WorkspaceApi, workdir: &WorkdirRegistryRecord, @@ -17037,10 +17203,7 @@ async fn scoped_create_external_workdir_grant( .ttl_seconds .map(|ttl_seconds| now + Duration::seconds(ttl_seconds as i64)) .map(|expires_at| expires_at.to_rfc3339_opts(SecondsFormat::Nanos, true)); - let permissions = match request.permission { - server_api::ExternalWorkdirPermission::ReadOnly => "read_only", - server_api::ExternalWorkdirPermission::ReadWrite => "read_write", - }; + let permissions = persist_external_workdir_permissions(request.permissions)?; let grant_id = format!("ewg-{}", Uuid::now_v7()); let workdir_id = format!("external-{}", Uuid::now_v7()); let grant = ExternalWorkdirGrantRecord { @@ -17080,27 +17243,24 @@ async fn scoped_create_external_workdir_grant( schedule_external_workdir_expiry(&api, &grant)?; Ok(( StatusCode::CREATED, - Json(external_workdir_grant_response(&grant)), + Json(external_workdir_grant_response(&grant)?), )) } fn external_workdir_grant_response( grant: &ExternalWorkdirGrantRecord, -) -> ExternalWorkdirGrantResponse { - ExternalWorkdirGrantResponse { +) -> Result { + Ok(ExternalWorkdirGrantResponse { grant_id: grant.grant_id.clone(), workspace_id: grant.workspace_id.clone(), working_directory_id: grant.workdir_id.clone(), provider_instance_id: grant.provider_instance_id.clone(), display_name: grant.display_name.clone(), - permission: match grant.permissions.as_str() { - "read_write" => server_api::ExternalWorkdirPermission::ReadWrite, - _ => server_api::ExternalWorkdirPermission::ReadOnly, - }, + permissions: external_workdir_permissions(&grant.permissions)?, expires_at: grant.expires_at.clone(), generation: grant.generation, status: grant.status.clone(), - } + }) } async fn scoped_external_workdir_provider_ws( @@ -17501,7 +17661,7 @@ async fn scoped_get_external_workdir_grant( cleanup_external_workdir_attachments(&api, &grant.workdir_id).await?; cancel_external_workdir_expiry(&api, &grant_id); } - Ok(Json(external_workdir_grant_response(&grant))) + Ok(Json(external_workdir_grant_response(&grant)?)) } async fn scoped_revoke_external_workdir_grant( @@ -17565,7 +17725,7 @@ async fn scoped_revoke_external_workdir_grant( // have committed its fence before a Runtime/session projection failure. cleanup_external_workdir_attachments(&api, &grant.workdir_id).await?; cancel_external_workdir_expiry(&api, &grant_id); - Ok(Json(external_workdir_grant_response(&grant))) + Ok(Json(external_workdir_grant_response(&grant)?)) } async fn cleanup_external_workdir_attachments( @@ -26068,6 +26228,39 @@ mod tests { WorkspaceRuntimeBinding, }; + #[test] + fn external_permission_categories_map_to_existing_capability_bits_without_escalation() { + assert_eq!( + external_workdir_capabilities("read_only").unwrap(), + workdir::WorkdirSessionCapabilities::READ_ONLY + ); + assert_eq!( + external_workdir_capabilities("read_write").unwrap(), + workdir::WorkdirSessionCapabilities::READ_WRITE + ); + assert_eq!( + external_workdir_capabilities("command_only").unwrap(), + workdir::WorkdirSessionCapabilities::COMMAND_ONLY + ); + assert_eq!( + external_workdir_capabilities("read_command").unwrap(), + workdir::WorkdirSessionCapabilities::READ_ONLY + .union(workdir::WorkdirSessionCapabilities::COMMAND_ONLY) + ); + assert_eq!( + external_workdir_capabilities("read_write_command").unwrap(), + workdir::WorkdirSessionCapabilities::ALL + ); + assert!( + persist_external_workdir_permissions(server_api::ExternalWorkdirPermissions { + read: false, + write: true, + command: false, + }) + .is_err() + ); + } + #[test] fn worker_projection_does_not_expose_unavailable_snapshot_as_current() { let availability = protocol::subscription::SubscriptionWorkerAvailability::Unavailable; @@ -26262,6 +26455,156 @@ mod tests { assert!(commands.try_recv().is_err()); } + #[tokio::test] + async fn external_operation_cancel_retries_without_blocking_a_saturated_provider_queue() { + let (sender, mut commands) = tokio::sync::mpsc::channel(1); + let queued_id = ExternalWorkdirOperationId::new("op-queued").unwrap(); + sender + .send(ExternalProviderCommand::Cancel { + operation_id: queued_id.clone(), + }) + .await + .unwrap(); + + let retried_id = ExternalWorkdirOperationId::new("op-retried").unwrap(); + enqueue_external_operation_cancel(sender, retried_id.clone()); + + let Some(ExternalProviderCommand::Cancel { operation_id }) = commands.recv().await else { + panic!("expected pre-existing queued cancellation"); + }; + assert_eq!(operation_id, queued_id); + let retried = tokio::time::timeout(std::time::Duration::from_secs(1), commands.recv()) + .await + .expect("cancellation retry should not be lost"); + let Some(ExternalProviderCommand::Cancel { operation_id }) = retried else { + panic!("expected retried cancellation"); + }; + assert_eq!(operation_id, retried_id); + } + + #[tokio::test] + async fn external_provider_command_only_session_denies_read_and_routes_bounded_lifecycle() { + let (sender, mut commands) = tokio::sync::mpsc::channel(4); + let capabilities = workdir::WorkdirSessionCapabilities::COMMAND_ONLY; + let connection = Arc::new(ExternalProviderConnection { + grant_id: "grant-command".to_string(), + workdir_id: "external-command".to_string(), + provider_instance_id: "provider-command".to_string(), + generation: 8, + expires_at: None, + capabilities, + admission: Arc::new(tokio::sync::Semaphore::new(16)), + sender, + }); + let session = Arc::new(ExternalProviderWorkdirSession::new(connection, None)); + let denied = session + .read(workdir::ReadRequest { + path: workdir::WorkdirPath::new("secret.txt").unwrap(), + offset: 0, + limit: 1, + max_bytes: 1, + }) + .await; + assert!(matches!( + denied, + Err(workdir::WorkdirError::Unsupported( + workdir::WorkdirSessionCapability::Read + )) + )); + assert!(commands.try_recv().is_err()); + + let starter = { + let session = session.clone(); + tokio::spawn(async move { + session + .start_command(workdir::CommandRequest { + command: "printf routed".to_string(), + timeout_secs: 30, + output_limit: 1024, + cwd: workdir::WorkdirPath::root(), + spill_dir: Some("/worker/private/spill".into()), + tool_call_id: Some("call-1".to_string()), + }) + .await + }) + }; + let Some(ExternalProviderCommand::Operation { + operation, + response, + .. + }) = commands.recv().await + else { + panic!("expected command start operation"); + }; + assert!(matches!( + operation, + WorkdirSessionOperation::CommandStart(workdir::CommandRequest { + spill_dir: None, + .. + }) + )); + let handle = workdir::CommandHandle("command-1".to_string()); + response + .send(Ok(WorkdirSessionOperationResult::CommandStart( + handle.clone(), + ))) + .unwrap(); + assert_eq!(starter.await.unwrap().unwrap(), handle); + + let closer = { + let session = session.clone(); + tokio::spawn(async move { session.close().await }) + }; + let Some(ExternalProviderCommand::Operation { + operation, + response, + .. + }) = commands.recv().await + else { + panic!("expected command cancel during close"); + }; + assert!(matches!( + operation, + WorkdirSessionOperation::CommandCancel(ref current) if current == &handle + )); + response + .send(Ok(WorkdirSessionOperationResult::CommandCancel)) + .unwrap(); + let Some(ExternalProviderCommand::Operation { + operation, + response, + .. + }) = commands.recv().await + else { + panic!("expected terminal command output during close"); + }; + assert!(matches!( + operation, + WorkdirSessionOperation::CommandOutput(workdir::CommandOutputRequest { + wait: true, + .. + }) + )); + 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, + }, + ))) + .unwrap(); + closer.await.unwrap().unwrap(); + assert!(matches!( + session.command_status(handle).await, + Err(workdir::WorkdirError::SessionClosed) + )); + } + #[tokio::test] async fn backend_expiry_fences_offline_grant_and_persists_cleanup_completion() { let dir = tempfile::tempdir().unwrap(); @@ -26696,7 +27039,7 @@ mod tests { } #[tokio::test] - async fn external_provider_websocket_reconnects_with_fresh_generation_and_revokes() { + async fn external_provider_websocket_routes_files_and_commands_then_reconnects_and_revokes() { let dir = tempfile::tempdir().unwrap(); init_clean_git_workspace(dir.path()); let api = test_api(dir.path()).await; @@ -26726,7 +27069,11 @@ mod tests { provider_instance_id: "provider-a".to_string(), display_name: "Session logs".to_string(), ttl_seconds: Some(600), - permission: server_api::ExternalWorkdirPermission::ReadWrite, + permissions: server_api::ExternalWorkdirPermissions { + read: true, + write: true, + command: true, + }, }), ) .await @@ -26755,7 +27102,7 @@ mod tests { .id() .clone(), generation, - capabilities: workdir::WorkdirSessionCapabilities::READ_WRITE, + capabilities: workdir::WorkdirSessionCapabilities::ALL, read_limits: workdir::BoundedReadLimits::EXTERNAL_DEFAULT, }, }) @@ -26867,14 +27214,16 @@ mod tests { ); assert_eq!( external_capabilities, - workdir::WorkdirSessionCapabilities::READ_WRITE + workdir::WorkdirSessionCapabilities::ALL ); let provider_root = tempfile::tempdir().unwrap(); fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").unwrap(); - let local_session = workdir::LocalWorkdirSession::external_read_write( + let pinned_provider_root = workdir::ExternalWorkdirRoot::pin(provider_root.path()).unwrap(); + let local_session = workdir::LocalWorkdirSession::external_with_capabilities_pinned( workdir::Workdir::new(&grant.working_directory_id), - provider_root.path(), + pinned_provider_root, workdir::BoundedReadLimits::EXTERNAL_DEFAULT, + workdir::WorkdirSessionCapabilities::ALL, ) .unwrap(); let session_lock = current_worker_session_lock(&api, &worker); @@ -26885,7 +27234,7 @@ mod tests { .unwrap(); assert_eq!( broker_session.capabilities(), - workdir::WorkdirSessionCapabilities::READ_WRITE + workdir::WorkdirSessionCapabilities::ALL ); drop(session_guard); let read = { @@ -27052,26 +27401,120 @@ mod tests { fs::read_to_string(provider_root.path().join("created.txt")).unwrap(), "created" ); - let command_error = broker_session - .start_command(workdir::CommandRequest { - command: "touch forbidden".to_string(), - timeout_secs: 1, - output_limit: 1024, - cwd: workdir::WorkdirPath::root(), - spill_dir: None, - tool_call_id: None, + let command_start = { + let broker_session = broker_session.clone(); + tokio::spawn(async move { + broker_session + .start_command(workdir::CommandRequest { + command: "printf command-ok".to_string(), + timeout_secs: 30, + output_limit: 1024, + cwd: workdir::WorkdirPath::root(), + spill_dir: Some("/worker/private/spill".into()), + tool_call_id: Some("call-e2e".to_string()), + }) + .await }) + }; + let (operation_id, operation) = loop { + let text = second.next().await.unwrap().unwrap().into_text().unwrap(); + let frame = serde_json::from_str::(&text).unwrap(); + if let ExternalWorkdirServerMessage::Operation { + generation: 2, + operation_id, + operation, + } = frame.message + { + break (operation_id, operation); + } + }; + assert!(matches!( + operation.as_inner(), + WorkdirSessionOperation::CommandStart(workdir::CommandRequest { + spill_dir: None, + .. + }) + )); + let result = workdir::http::dispatch_workdir_session_operation( + &local_session, + operation.into_inner(), + ) + .await + .unwrap() + .try_into() + .unwrap(); + second + .send(Message::Text( + serde_json::to_string(&ExternalWorkdirProviderFrame::current( + ExternalWorkdirProviderMessage::OperationResult { + generation: 2, + operation_id, + outcome: ExternalWorkdirOperationOutcome::Completed { result }, + }, + )) + .unwrap() + .into(), + )) .await - .unwrap_err(); - assert!( - matches!( - command_error, - workdir::WorkdirError::Unsupported(workdir::WorkdirSessionCapability::Command) - | workdir::WorkdirError::UnsupportedOperation(_) - | workdir::WorkdirError::Denied(_) - ), - "unexpected command error: {command_error:?}" - ); + .unwrap(); + let command_handle = command_start.await.unwrap().unwrap(); + + let command_output = { + let broker_session = broker_session.clone(); + let handle = command_handle.clone(); + tokio::spawn(async move { + broker_session + .command_output(workdir::CommandOutputRequest { + handle, + cursor: 0, + limit: 1024, + wait: true, + }) + .await + }) + }; + let (operation_id, operation) = loop { + let text = second.next().await.unwrap().unwrap().into_text().unwrap(); + let frame = serde_json::from_str::(&text).unwrap(); + if let ExternalWorkdirServerMessage::Operation { + generation: 2, + operation_id, + operation, + } = frame.message + { + break (operation_id, operation); + } + }; + assert!(matches!( + operation.as_inner(), + WorkdirSessionOperation::CommandOutput(_) + )); + let result = workdir::http::dispatch_workdir_session_operation( + &local_session, + operation.into_inner(), + ) + .await + .unwrap() + .try_into() + .unwrap(); + second + .send(Message::Text( + serde_json::to_string(&ExternalWorkdirProviderFrame::current( + ExternalWorkdirProviderMessage::OperationResult { + generation: 2, + operation_id, + outcome: ExternalWorkdirOperationOutcome::Completed { result }, + }, + )) + .unwrap() + .into(), + )) + .await + .unwrap(); + let command_output = command_output.await.unwrap().unwrap(); + assert_eq!(command_output.content, "command-ok"); + assert_eq!(command_output.exit_code, Some(0)); + assert_eq!(command_output.output_path, None); let Json(revoked) = scoped_revoke_external_workdir_grant( State(api.clone()), diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index da4eaaa37..9c902f5b3 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -21,7 +21,7 @@ use crate::workspace_deletion::WorkspaceDeletionStore; use crate::{Error, Result}; const OLDEST_SCHEMA_VERSION: i64 = 50; -const LATEST_SCHEMA_VERSION: i64 = 67; +const LATEST_SCHEMA_VERSION: i64 = 68; const SCHEMA_BASELINE_NAME: &str = "workspace schema baseline"; const WORKSPACE_RUNTIME_BINDINGS_MIGRATION_NAME: &str = "workspace runtime bindings"; const RUNTIME_BINDING_AUDIT_MIGRATION_NAME: &str = "workspace Runtime binding revision and audit"; @@ -50,6 +50,8 @@ const TICKET_TARGETS_AND_WORKDIR_CAPABILITIES_MIGRATION_NAME: &str = "Ticket target collection and Workdir attachment capabilities"; const EXTERNAL_WORKDIR_ACCESS_AND_OPTIONAL_EXPIRY_MIGRATION_NAME: &str = "External Workdir read-write access and optional expiry"; +const EXTERNAL_WORKDIR_COMMAND_PERMISSION_MIGRATION_NAME: &str = + "External Workdir explicit command permission"; const TICKET_SCHEMA_VERSION_WITH_TARGETS: i64 = 7; const TICKET_SCHEMA_BASELINE_NAME: &str = "ticket schema baseline"; const WORKER_REGISTRY_PROJECTION_SCHEMA: &str = r#" @@ -182,6 +184,11 @@ const MIGRATIONS: &[Migration] = &[ name: EXTERNAL_WORKDIR_ACCESS_AND_OPTIONAL_EXPIRY_MIGRATION_NAME, apply: migrate_external_workdir_access_and_optional_expiry_v66_to_v67, }, + Migration { + version: 68, + name: EXTERNAL_WORKDIR_COMMAND_PERMISSION_MIGRATION_NAME, + apply: migrate_external_workdir_command_permission_v67_to_v68, + }, ]; #[derive(Clone, Copy)] @@ -12098,6 +12105,90 @@ fn migrate_external_workdir_access_and_optional_expiry_v66_to_v67(conn: &Connect } } +fn migrate_external_workdir_command_permission_v67_to_v68(conn: &Connection) -> Result<()> { + let current = current_schema_version(conn)?; + if current != 67 { + return Err(Error::Store(format!( + "expected schema version 67 before {EXTERNAL_WORKDIR_COMMAND_PERMISSION_MIGRATION_NAME} migration, found {current}" + ))); + } + let foreign_keys_enabled = + conn.query_row("PRAGMA foreign_keys", [], |row| row.get::<_, i64>(0))? != 0; + if foreign_keys_enabled { + conn.execute_batch("PRAGMA foreign_keys = OFF;")?; + } + let result = (|| { + let tx = rusqlite::Transaction::new_unchecked(conn, TransactionBehavior::Exclusive)?; + tx.execute_batch( + r#" + CREATE TABLE external_workdir_grants_v68 ( + grant_id TEXT NOT NULL, + workspace_id TEXT NOT NULL, + workdir_id TEXT NOT NULL, + provider_instance_id TEXT NOT NULL, + display_name TEXT NOT NULL, + permissions TEXT NOT NULL CHECK (permissions IN ( + 'read_only', 'read_write', 'command_only', + 'read_command', 'read_write_command' + )), + created_by TEXT NOT NULL, + created_at TEXT NOT NULL, + expires_at TEXT, + generation INTEGER NOT NULL CHECK (generation > 0), + status TEXT NOT NULL CHECK (status IN ('pending', 'online', 'offline', 'revoked', 'expired')), + updated_at TEXT NOT NULL, + cleanup_state TEXT NOT NULL DEFAULT 'not_required' CHECK (cleanup_state IN ('not_required', 'pending', 'retry_required', 'completed')), + cleanup_error TEXT, + cleanup_updated_at TEXT, + PRIMARY KEY (workspace_id, grant_id), + UNIQUE (workspace_id, workdir_id), + FOREIGN KEY (workspace_id) REFERENCES workspaces(workspace_id) ON DELETE CASCADE, + FOREIGN KEY (workspace_id, workdir_id) + REFERENCES workdir_registry(workspace_id, workdir_id) ON DELETE CASCADE DEFERRABLE INITIALLY DEFERRED + ); + INSERT INTO external_workdir_grants_v68 ( + grant_id, workspace_id, workdir_id, provider_instance_id, display_name, + permissions, created_by, created_at, expires_at, generation, status, updated_at, + cleanup_state, cleanup_error, cleanup_updated_at + ) SELECT grant_id, workspace_id, workdir_id, provider_instance_id, display_name, + permissions, created_by, created_at, expires_at, generation, status, updated_at, + cleanup_state, cleanup_error, cleanup_updated_at + FROM external_workdir_grants; + DROP TABLE external_workdir_grants; + ALTER TABLE external_workdir_grants_v68 RENAME TO external_workdir_grants; + CREATE INDEX idx_external_workdir_grants_status_expiry + ON external_workdir_grants(workspace_id, status, expires_at); + "#, + )?; + let violations = + tx.query_row("SELECT COUNT(*) FROM pragma_foreign_key_check", [], |row| { + row.get::<_, i64>(0) + })?; + if violations != 0 { + return Err(Error::Store(format!( + "schema-68 migration left {violations} foreign-key violation(s)" + ))); + } + tx.execute( + "INSERT INTO __yoi_schema_migrations (version, name) VALUES (?1, ?2)", + params![68_i64, EXTERNAL_WORKDIR_COMMAND_PERMISSION_MIGRATION_NAME], + )?; + tx.commit()?; + Ok(()) + })(); + let restore = if foreign_keys_enabled { + conn.execute_batch("PRAGMA foreign_keys = ON;") + .map_err(Error::from) + } else { + Ok(()) + }; + match (result, restore) { + (Err(error), _) => Err(error), + (Ok(()), Err(error)) => Err(error), + (Ok(()), Ok(())) => Ok(()), + } +} + fn create_latest_workspace_schema(conn: &Connection) -> Result<()> { conn.execute_batch("DROP TABLE IF EXISTS typed_ticket_targets;")?; conn.execute_batch(include_str!("latest_schema.sql"))?; @@ -14233,6 +14324,10 @@ mod tests { version: 67, name: EXTERNAL_WORKDIR_ACCESS_AND_OPTIONAL_EXPIRY_MIGRATION_NAME.to_string(), }, + WorkspaceSchemaMigrationStep { + version: 68, + name: EXTERNAL_WORKDIR_COMMAND_PERMISSION_MIGRATION_NAME.to_string(), + }, ] ); @@ -14295,6 +14390,10 @@ mod tests { 67, EXTERNAL_WORKDIR_ACCESS_AND_OPTIONAL_EXPIRY_MIGRATION_NAME.to_string(), ), + ( + 68, + EXTERNAL_WORKDIR_COMMAND_PERMISSION_MIGRATION_NAME.to_string(), + ), ] ); assert!(!table_exists(conn, "trusted_runtime_records")?); @@ -14348,6 +14447,68 @@ mod tests { .unwrap(); } + #[test] + fn schema_v68_preserves_legacy_external_grants_without_command_authority() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("server.db"); + prepare_schema_v50(&path, Some("workspace-a")); + let mut conn = Connection::open(&path).unwrap(); + configure_sqlite(&conn).unwrap(); + for migration in MIGRATIONS + .iter() + .filter(|migration| migration.version <= 67) + { + (migration.apply)(&conn).unwrap(); + } + let tx = conn.transaction().unwrap(); + tx.execute_batch( + "INSERT INTO external_workdir_grants ( + grant_id, workspace_id, workdir_id, provider_instance_id, display_name, + permissions, created_by, created_at, expires_at, generation, status, updated_at + ) VALUES + ('grant-ro', 'workspace-a', 'external-ro', 'provider-ro', 'Read only', + 'read_only', 'owner', '1', NULL, 1, 'pending', '1'), + ('grant-rw', 'workspace-a', 'external-rw', 'provider-rw', 'Read write', + 'read_write', 'owner', '1', NULL, 1, 'pending', '1'); + INSERT INTO workdir_registry ( + workspace_id, workdir_id, display_name, source_kind, runtime_id, + repository_id, external_grant_id, creation_selector, creation_ref, + materialization_status, cleanliness, created_at, updated_at + ) VALUES + ('workspace-a', 'external-ro', 'Read only', 'external_grant', NULL, + NULL, 'grant-ro', NULL, NULL, 'pending', 'unknown', '1', '1'), + ('workspace-a', 'external-rw', 'Read write', 'external_grant', NULL, + NULL, 'grant-rw', NULL, NULL, 'pending', 'unknown', '1', '1');", + ) + .unwrap(); + tx.commit().unwrap(); + + migrate_external_workdir_command_permission_v67_to_v68(&conn).unwrap(); + assert_eq!(current_schema_version(&conn).unwrap(), 68); + let permissions = conn + .prepare("SELECT permissions FROM external_workdir_grants ORDER BY grant_id") + .unwrap() + .query_map([], |row| row.get::<_, String>(0)) + .unwrap() + .collect::>>() + .unwrap(); + assert_eq!(permissions, vec!["read_only", "read_write"]); + assert!( + conn.execute( + "UPDATE external_workdir_grants SET permissions = 'read_command' WHERE grant_id = 'grant-ro'", + [], + ) + .is_ok() + ); + assert!( + conn.execute( + "UPDATE external_workdir_grants SET permissions = 'write_only' WHERE grant_id = 'grant-rw'", + [], + ) + .is_err() + ); + } + #[test] fn migration_resumes_from_a_valid_partially_applied_chain() { let temp = tempfile::tempdir().unwrap(); @@ -14366,7 +14527,7 @@ mod tests { .map(|migration| migration.version) .collect::>(), vec![ - 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67 + 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68 ] ); SqliteWorkspaceStore::migrate_database(&path).unwrap(); @@ -14375,7 +14536,7 @@ mod tests { current_schema_version(&conn).unwrap(), LATEST_SCHEMA_VERSION ); - assert_eq!(workspace_schema_migration_history(&conn).unwrap().len(), 18); + assert_eq!(workspace_schema_migration_history(&conn).unwrap().len(), 19); assert!(column_exists(&conn, "worker_workdir_links", "alias").unwrap()); assert!(column_exists(&conn, "worker_workdir_links", "capabilities").unwrap()); assert!(!column_exists(&conn, "worker_workdir_links", "role").unwrap()); diff --git a/crates/yoi/src/main.rs b/crates/yoi/src/main.rs index e0c0ce435..7a3577091 100644 --- a/crates/yoi/src/main.rs +++ b/crates/yoi/src/main.rs @@ -623,6 +623,7 @@ fn parse_workdir_args( let mut ttl = None; let mut read_only = false; let mut read_write = false; + let mut explicit_permissions: Option = None; let mut non_interactive = false; let mut index = 0; while index < args.len() { @@ -666,6 +667,32 @@ fn parse_workdir_args( ); index += 2; } + "--permission" => { + let category = required_option_value(args, index, "--permission")?; + let permissions = + explicit_permissions.get_or_insert(server_api::ExternalWorkdirPermissions { + read: false, + write: false, + command: false, + }); + let enabled = match category.to_ascii_lowercase().as_str() { + "read" => &mut permissions.read, + "write" => &mut permissions.write, + "command" => &mut permissions.command, + _ => { + return Err(ParseError(format!( + "unknown --permission category `{category}`; expected read, write, or command" + ))); + } + }; + if *enabled { + return Err(ParseError(format!( + "--permission {category} may only be provided once" + ))); + } + *enabled = true; + index += 2; + } "--read-only" => { if read_only { return Err(ParseError( @@ -715,6 +742,21 @@ fn parse_workdir_args( "--read-only conflicts with --read-write".to_string(), )); } + if explicit_permissions.is_some() && (read_only || read_write) { + return Err(ParseError( + "--permission conflicts with legacy --read-only/--read-write".to_string(), + )); + } + let permissions = explicit_permissions.unwrap_or(server_api::ExternalWorkdirPermissions { + read: true, + write: read_write, + command: false, + }); + if permissions.write && !permissions.read { + return Err(ParseError( + "--permission write requires --permission read for read-before-write".to_string(), + )); + } if non_interactive && workspace_id.is_none() { return Err(ParseError( "--non-interactive requires --workspace-id".to_string(), @@ -728,11 +770,7 @@ fn parse_workdir_args( backend_url, display_name: display_name.unwrap_or_else(|| "External Workdir".to_string()), ttl, - permission: if read_write { - server_api::ExternalWorkdirPermission::ReadWrite - } else { - server_api::ExternalWorkdirPermission::ReadOnly - }, + permissions, non_interactive, })) } @@ -1862,7 +1900,7 @@ Usage: yoi --backend [--workspace-id ] panel yoi [--backend ] init --display-name --repository-key [--repository ] [--default-ref ] yoi [--backend ] login [--no-wait] - yoi [--backend ] workdir share [--workspace-id ] [--read-only|--read-write] [--ttl ] [--display-name ] [--non-interactive] + yoi [--backend ] workdir share [--workspace-id ] [--permission ]... [--read-only|--read-write] [--ttl ] [--display-name ] [--non-interactive] yoi [OPTIONS] Target selection: @@ -1908,7 +1946,7 @@ Options: fn print_workdir_help() { println!( - "yoi workdir share\n\nUsage:\n yoi [--backend ] workdir share [OPTIONS]\n\nThe foreground CLI provides one local directory over an authenticated outbound Backend connection.\nInteractive mode is the default: choose an accessible Workspace, permission, and TTL, then explicitly confirm before any grant is created.\nRead-only is the default permission. --read-write enables the existing filesystem Write/Edit operations only; Bash, command, and exec remain disabled.\nOmitting --ttl means unlimited automatic lifetime, but stop, disconnect, Ctrl-C, and revoke still fail closed. TTL accepts s, m, or h suffixes from 60s through 24h.\nNo host path is sent to the Backend. Ctrl-C explicitly revokes the grant before exit.\n\nOptions:\n --workspace-id Preselect a Workspace; required with --non-interactive\n --read-only Explicitly select the default read-only permission\n --read-write Allow filesystem Write/Edit operations (never commands)\n --ttl Automatic expiry; omission is unlimited\n --display-name Backend-visible label (default: External Workdir)\n --non-interactive Never prompt; requires --workspace-id\n -h, --help Print help\n" + "yoi workdir share\n\nUsage:\n yoi [--backend ] workdir share [OPTIONS]\n\nThe foreground CLI provides one local directory over an authenticated outbound Backend connection.\nInteractive mode is the default: choose an accessible Workspace, READ/WRITE/COMMAND categories, and TTL, then explicitly confirm before any grant is created.\nREAD enables Read/Glob/Grep and related observation operations. WRITE enables Write/Edit and requires READ for read-before-write. COMMAND is independent, disabled by default, and runs with the sharing CLI user's host authority; it is not sandboxed.\nRepeat --permission to choose explicit categories in automation. Legacy --read-only and --read-write remain file-only aliases and conflict with --permission.\nOmitting --ttl means unlimited automatic lifetime, but stop, disconnect, Ctrl-C, and revoke still fail closed. TTL accepts s, m, or h suffixes from 60s through 24h.\nNo host path is sent to the Backend. Ctrl-C explicitly revokes the grant before exit.\n\nOptions:\n --workspace-id Preselect a Workspace; required with --non-interactive\n --permission Enable read, write, or command; repeat per category\n --read-only Legacy alias for --permission read\n --read-write Legacy alias for read + write (commands remain disabled)\n --ttl Automatic expiry; omission is unlimited\n --display-name Backend-visible label (default: External Workdir)\n --non-interactive Never prompt; requires --workspace-id\n -h, --help Print help\n" ); } @@ -1989,12 +2027,43 @@ mod tests { assert_eq!(options.display_name, "session-analysis"); assert_eq!(options.ttl, Some(Duration::from_secs(5_400))); assert_eq!( - options.permission, - server_api::ExternalWorkdirPermission::ReadWrite + options.permissions, + server_api::ExternalWorkdirPermissions { + read: true, + write: true, + command: false, + } ); assert!(options.non_interactive); } + #[test] + fn workdir_share_parses_independent_command_permission() { + let args = vec![ + "share".to_string(), + ".".to_string(), + "--permission".to_string(), + "command".to_string(), + "--non-interactive".to_string(), + ]; + let target = TargetSelection { + backend_url: Some("https://backend.example".to_string()), + workspace_id: Some("workspace-a".to_string()), + ..TargetSelection::default() + }; + let Mode::WorkdirShare(options) = parse_workdir_args(&args, &target).unwrap() else { + panic!("expected Workdir share mode"); + }; + assert_eq!( + options.permissions, + server_api::ExternalWorkdirPermissions { + read: false, + write: false, + command: true, + } + ); + } + #[test] fn workdir_share_defaults_to_interactive_read_only_and_unlimited() { let args = vec!["share".to_string(), ".".to_string()]; @@ -2009,8 +2078,8 @@ mod tests { assert_eq!(options.display_name, "External Workdir"); assert_eq!(options.ttl, None); assert_eq!( - options.permission, - server_api::ExternalWorkdirPermission::ReadOnly + options.permissions, + server_api::ExternalWorkdirPermissions::default() ); assert!(!options.non_interactive); } @@ -2048,6 +2117,33 @@ mod tests { .contains("conflicts") ); + let legacy_conflict = vec![ + "share".to_string(), + ".".to_string(), + "--read-write".to_string(), + "--permission".to_string(), + "command".to_string(), + ]; + assert!( + parse_workdir_args(&legacy_conflict, &target) + .unwrap_err() + .0 + .contains("conflicts") + ); + + let write_without_read = vec![ + "share".to_string(), + ".".to_string(), + "--permission".to_string(), + "write".to_string(), + ]; + assert!( + parse_workdir_args(&write_without_read, &target) + .unwrap_err() + .0 + .contains("requires --permission read") + ); + let missing_workspace = TargetSelection { backend_url: Some("https://backend.example".to_string()), ..TargetSelection::default() diff --git a/crates/yoi/src/workdir_share.rs b/crates/yoi/src/workdir_share.rs index 44d104abf..82d74ec8b 100644 --- a/crates/yoi/src/workdir_share.rs +++ b/crates/yoi/src/workdir_share.rs @@ -30,12 +30,23 @@ pub(crate) struct WorkdirShareOptions { pub backend_url: String, pub display_name: String, pub ttl: Option, - pub permission: server_api::ExternalWorkdirPermission, + pub permissions: server_api::ExternalWorkdirPermissions, pub non_interactive: bool, } const PROVIDER_OPERATION_TIMEOUT: Duration = Duration::from_secs(30); +fn provider_operation_timeout(operation: &workdir::http::WorkdirSessionOperation) -> Duration { + if matches!( + operation, + workdir::http::WorkdirSessionOperation::CommandOutput(request) if request.wait + ) { + Duration::from_secs(workdir::external::MAX_EXTERNAL_COMMAND_TIMEOUT_SECS + 10) + } else { + PROVIDER_OPERATION_TIMEOUT + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ProviderEnd { Disconnected, @@ -128,13 +139,47 @@ fn authenticated_server_api_client( Ok(ServerApiClient::with_client(base_url, http_client)) } -fn permission_label(permission: server_api::ExternalWorkdirPermission) -> &'static str { - match permission { - server_api::ExternalWorkdirPermission::ReadOnly => "read-only", - server_api::ExternalWorkdirPermission::ReadWrite => { - "read-write (files only; commands disabled)" - } +fn permissions_label(permissions: server_api::ExternalWorkdirPermissions) -> String { + let mut enabled = Vec::new(); + if permissions.read { + enabled.push("READ"); } + if permissions.write { + enabled.push("WRITE"); + } + if permissions.command { + enabled.push("COMMAND"); + } + enabled.join(", ") +} + +fn validate_permissions(permissions: server_api::ExternalWorkdirPermissions) -> Result<(), String> { + if permissions.write && !permissions.read { + return Err( + "WRITE requires READ so existing files can satisfy read-before-write; select both categories" + .to_string(), + ); + } + if !permissions.read && !permissions.command { + return Err("select at least one of READ or COMMAND".to_string()); + } + Ok(()) +} + +fn workdir_capabilities( + permissions: server_api::ExternalWorkdirPermissions, +) -> workdir::WorkdirSessionCapabilities { + let mut capabilities = workdir::WorkdirSessionCapabilities::EMPTY; + if permissions.read { + capabilities = capabilities.union(workdir::WorkdirSessionCapabilities::READ_ONLY); + } + if permissions.write { + capabilities = capabilities.union(workdir::WorkdirSessionCapabilities::WRITE_ONLY); + } + if permissions.command { + capabilities = capabilities.union(workdir::WorkdirSessionCapabilities::COMMAND_ONLY); + } + capabilities } #[derive(Debug, Clone, PartialEq, Eq)] @@ -144,7 +189,7 @@ struct ResolvedWorkdirShareOptions { backend_url: String, display_name: String, ttl: Option, - permission: server_api::ExternalWorkdirPermission, + permissions: server_api::ExternalWorkdirPermissions, } pub(crate) fn parse_ttl(value: &str) -> Result { @@ -254,30 +299,49 @@ fn configure_interactively( }; options.workspace_id = Some(workspaces[selected_index].workspace_id.clone()); - let current_permission = match options.permission { - server_api::ExternalWorkdirPermission::ReadOnly => "r", - server_api::ExternalWorkdirPermission::ReadWrite => "w", - }; writeln!( output, - "Permission: read-only or read-write files. Read-write never enables Bash/command/exec." + "Permissions are independent categories: READ (Read/Glob/Grep), WRITE (Write/Edit), and COMMAND (host command lifecycle). WRITE requires READ for read-before-write. COMMAND runs with this CLI user's host authority and is not sandboxed." ) .map_err(|error| format!("failed to write prompt: {error}"))?; - let Some(permission) = read_prompt_line( + let Some(permission_input) = read_prompt_line( input, output, - &format!("Permission [r/w; current {current_permission}; q to cancel]: "), + &format!( + "Permissions [comma-separated READ/WRITE/COMMAND; current {}; q to cancel]: ", + permissions_label(options.permissions) + ), )? else { return Ok(None); }; - options.permission = match permission.as_str() { - "" => options.permission, - "r" | "read-only" | "read_only" => server_api::ExternalWorkdirPermission::ReadOnly, - "w" | "read-write" | "read_write" => server_api::ExternalWorkdirPermission::ReadWrite, - "q" | "quit" | "cancel" => return Ok(None), - _ => return Err("Permission must be read-only (`r`) or read-write (`w`)".to_string()), - }; + if matches!(permission_input.as_str(), "q" | "quit" | "cancel") { + return Ok(None); + } + if !permission_input.is_empty() { + let mut permissions = server_api::ExternalWorkdirPermissions { + read: false, + write: false, + command: false, + }; + for category in permission_input + .split([',', ' ']) + .filter(|category| !category.is_empty()) + { + match category.to_ascii_lowercase().as_str() { + "r" | "read" | "read-only" | "read_only" => permissions.read = true, + "w" | "write" | "read-write" | "read_write" => permissions.write = true, + "c" | "command" => permissions.command = true, + _ => { + return Err(format!( + "unknown permission category `{category}`; use READ, WRITE, or COMMAND" + )); + } + } + } + validate_permissions(permissions)?; + options.permissions = permissions; + } writeln!(output, "TTL: {}", ttl_label(options.ttl)) .map_err(|error| format!("failed to write prompt: {error}"))?; @@ -305,8 +369,8 @@ fn configure_interactively( .and_then(|()| { writeln!( output, - " Permission: {}", - permission_label(options.permission) + " Permissions: {}", + permissions_label(options.permissions) ) }) .and_then(|()| writeln!(output, " TTL: {}", ttl_label(options.ttl))) @@ -324,7 +388,7 @@ fn configure_interactively( backend_url: options.backend_url, display_name: options.display_name, ttl: options.ttl, - permission: options.permission, + permissions: options.permissions, })) } @@ -342,6 +406,7 @@ fn ensure_interactive_tty(stdin_is_tty: bool, stderr_is_tty: bool) -> Result<(), async fn resolve_share_options( options: WorkdirShareOptions, ) -> Result, String> { + validate_permissions(options.permissions)?; if options.non_interactive { let workspace_id = options.workspace_id.clone().ok_or_else(|| { "non-interactive sharing requires --workspace-id; no Workspace is inferred from cwd or path" @@ -353,7 +418,7 @@ async fn resolve_share_options( backend_url: options.backend_url, display_name: options.display_name, ttl: options.ttl, - permission: options.permission, + permissions: options.permissions, })); } ensure_interactive_tty( @@ -392,35 +457,29 @@ pub(crate) async fn run(options: WorkdirShareOptions) -> Result<(), String> { provider_instance_id: provider_instance_id.clone(), display_name: options.display_name.clone(), ttl_seconds: options.ttl.map(|ttl| ttl.as_secs()), - permission: options.permission, + permissions: options.permissions, }, ) .await .map_err(|error| format!("External Workdir grant request failed: {error}"))?; - let limits = BoundedReadLimits::EXTERNAL_DEFAULT; - let session = match options.permission { - server_api::ExternalWorkdirPermission::ReadOnly => { - LocalWorkdirSession::external_read_only_pinned( - Workdir::new(&grant.working_directory_id), - pinned_root, - limits, - ) - } - server_api::ExternalWorkdirPermission::ReadWrite => { - LocalWorkdirSession::external_read_write_pinned( - Workdir::new(&grant.working_directory_id), - pinned_root, - limits, - ) - } + if grant.permissions != options.permissions { + return Err( + "Backend returned permissions that differ from the requested grant".to_string(), + ); } - .map_err(|error| error.to_string())?; + let capabilities = workdir_capabilities(grant.permissions); + let limits = BoundedReadLimits::EXTERNAL_DEFAULT; println!("External Workdir grant: {}", grant.grant_id); println!("Logical Workdir: {}", grant.working_directory_id); println!("Display name: {}", grant.display_name); - println!("Permission: {}", permission_label(grant.permission)); + println!("Permissions: {}", permissions_label(grant.permissions)); + if grant.permissions.command { + println!( + "COMMAND executes with this sharing CLI user's host authority; it is not sandboxed." + ); + } println!( "Expires at: {}", grant @@ -440,6 +499,16 @@ pub(crate) async fn run(options: WorkdirShareOptions) -> Result<(), String> { .transpose()?; let mut generation = grant.generation; loop { + // Command handles and results are generation-owned. A reconnect always + // gets a fresh LocalWorkdirSession, and closing the old one kills and + // reaps every process before the next generation can register. + let session = LocalWorkdirSession::external_with_capabilities_pinned( + Workdir::new(&grant.working_directory_id), + pinned_root.clone(), + limits, + capabilities, + ) + .map_err(|error| error.to_string())?; let outcome = serve_provider_connection( &backend_client, &server_api_client, @@ -450,6 +519,17 @@ pub(crate) async fn run(options: WorkdirShareOptions) -> Result<(), String> { generation, ) .await; + let close_result = session.close().await.map_err(|error| error.to_string()); + let outcome = match (outcome, close_result) { + (Ok(outcome), Ok(())) => Ok(outcome), + (Err(error), Ok(())) => Err(error), + (Ok(_), Err(error)) => Err(format!( + "External Workdir provider session cleanup failed: {error}" + )), + (Err(connection_error), Err(close_error)) => Err(format!( + "{connection_error}; provider session cleanup failed: {close_error}" + )), + }; let outcome = match outcome { Ok(outcome) => outcome, Err(connection_error) => { @@ -625,7 +705,7 @@ async fn serve_provider_connection( tokio::select! { signal = tokio::signal::ctrl_c() => { signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?; - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?; return Ok(ProviderEnd::Interrupted); } @@ -640,11 +720,12 @@ async fn serve_provider_connection( let cancelled = match operation.join() { Ok(cancelled) => cancelled, Err(error) => { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Err(error); } }; let outcome = if cancelled { + cancel_completed_command(session, &completion.outcome).await; ExternalWorkdirOperationOutcome::Cancelled } else { completion.outcome @@ -657,19 +738,19 @@ async fn serve_provider_connection( }, ); if send_provider_frame(&mut socket, &result).await.is_err() { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Ok(ProviderEnd::Disconnected); } } incoming = socket.next() => { let Some(incoming) = incoming else { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Ok(ProviderEnd::Disconnected); }; let incoming = match incoming { Ok(incoming) => incoming, Err(_) => { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Ok(ProviderEnd::Disconnected); } }; @@ -678,7 +759,7 @@ async fn serve_provider_connection( let frame = match serde_json::from_str::(&text) { Ok(frame) => frame, Err(error) => { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Err(format!("Backend returned an invalid External Workdir frame: {error}")); } }; @@ -689,24 +770,26 @@ async fn serve_provider_connection( operation, } if frame_generation == generation => { if operations.len() >= 16 || operations.contains_key(operation_id.as_str()) { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Err("Backend exceeded External Workdir operation bounds".to_string()); } let key = operation_id.as_str().to_string(); let completion_sender = completion_sender.clone(); let task_operation_id = operation_id.clone(); let cancelled = Arc::new(AtomicBool::new(false)); + let operation_timeout = provider_operation_timeout(operation.as_inner()); let guarded_session = session.with_operation_guard( cancelled.clone(), - Instant::now() + PROVIDER_OPERATION_TIMEOUT, + Instant::now() + operation_timeout, ); let operation = operation.into_inner(); + let runtime = tokio::runtime::Handle::current(); let spawned = std::thread::Builder::new() .name("external-workdir-operation".to_string()) .spawn(move || { let outcome = std::panic::catch_unwind( std::panic::AssertUnwindSafe(|| { - futures::executor::block_on(execute_operation( + runtime.block_on(execute_operation( &guarded_session, operation, )) @@ -726,7 +809,7 @@ async fn serve_provider_connection( let thread = match spawned { Ok(thread) => thread, Err(error) => { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Err(format!("failed to start External Workdir operation: {error}")); } }; @@ -746,7 +829,7 @@ async fn serve_provider_connection( ExternalWorkdirProviderMessage::Heartbeat { generation, sequence }, ); if send_provider_frame(&mut socket, &heartbeat).await.is_err() { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Ok(ProviderEnd::Disconnected); } } @@ -764,7 +847,7 @@ async fn serve_provider_connection( generation: frame_generation, .. } if frame_generation == generation => { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; let acknowledged = ExternalWorkdirProviderFrame::current( ExternalWorkdirProviderMessage::RevokeAcknowledged { generation }, ); @@ -772,7 +855,7 @@ async fn serve_provider_connection( return Ok(ProviderEnd::Revoked); } _ => { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Err("stale or unexpected External Workdir provider frame".to_string()); } } @@ -784,12 +867,12 @@ async fn serve_provider_connection( ) .await; if !matches!(pong, Ok(Ok(()))) { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Ok(ProviderEnd::Disconnected); } } Message::Close(_) => { - stop_operations(&mut operations, &mut completions).await?; + stop_provider_operations(session, &mut operations, &mut completions).await?; return Ok(ProviderEnd::Disconnected); } _ => {} @@ -799,6 +882,28 @@ async fn serve_provider_connection( } } +async fn cancel_completed_command( + session: &LocalWorkdirSession, + outcome: &ExternalWorkdirOperationOutcome, +) { + let ExternalWorkdirOperationOutcome::Completed { result } = outcome else { + return; + }; + let workdir::http::WorkdirSessionOperationResult::CommandStart(handle) = result.as_inner() + else { + return; + }; + let _ = session.cancel_command(handle.clone()).await; + let _ = session + .command_output(workdir::CommandOutputRequest { + handle: handle.clone(), + cursor: 0, + limit: workdir::external::MAX_EXTERNAL_COMMAND_OUTPUT_BYTES, + wait: true, + }) + .await; +} + async fn execute_operation( session: &LocalWorkdirSession, operation: workdir::http::WorkdirSessionOperation, @@ -825,6 +930,15 @@ fn cancel_operations(operations: &HashMap) { } } +async fn stop_provider_operations( + session: &LocalWorkdirSession, + operations: &mut HashMap, + completions: &mut tokio::sync::mpsc::Receiver, +) -> Result<(), String> { + session.close().await.map_err(|error| error.to_string())?; + stop_operations(operations, completions).await +} + async fn stop_operations( operations: &mut HashMap, completions: &mut tokio::sync::mpsc::Receiver, @@ -925,7 +1039,7 @@ mod tests { backend_url: "https://backend.example".to_string(), display_name: "External Workdir".to_string(), ttl: None, - permission: server_api::ExternalWorkdirPermission::ReadOnly, + permissions: server_api::ExternalWorkdirPermissions::default(), non_interactive: false, } } @@ -954,8 +1068,8 @@ mod tests { assert_eq!(selected.workspace_id, "workspace-a"); assert_eq!( - selected.permission, - server_api::ExternalWorkdirPermission::ReadOnly + selected.permissions, + server_api::ExternalWorkdirPermissions::default() ); assert_eq!(selected.ttl, None); assert!( @@ -966,12 +1080,12 @@ mod tests { } #[test] - fn interactive_share_requires_explicit_read_write_selection_and_confirmation() { + fn interactive_share_requires_explicit_write_and_command_selection_and_confirmation() { let workspaces = [ workspace("workspace-a", "Alpha"), workspace("workspace-b", "Beta"), ]; - let mut input = std::io::Cursor::new(b"2\nw\n15m\ny\n".to_vec()); + let mut input = std::io::Cursor::new(b"2\nREAD,WRITE,COMMAND\n15m\ny\n".to_vec()); let mut output = Vec::new(); let selected = @@ -981,10 +1095,17 @@ mod tests { assert_eq!(selected.workspace_id, "workspace-b"); assert_eq!( - selected.permission, - server_api::ExternalWorkdirPermission::ReadWrite + selected.permissions, + server_api::ExternalWorkdirPermissions { + read: true, + write: true, + command: true, + } ); assert_eq!(selected.ttl, Some(Duration::from_secs(900))); + let output = String::from_utf8(output).unwrap(); + assert!(output.contains("not sandboxed")); + assert!(output.contains("Permissions: READ, WRITE, COMMAND")); } #[test] @@ -1101,7 +1222,7 @@ mod tests { "working_directory_id": "workdir-a", "provider_instance_id": "provider-a", "display_name": "Shared files", - "permission": "read_only", + "permissions": {"read": true, "write": false, "command": false}, "expires_at": "2026-09-23T03:00:00Z", "generation": 7, "status": "pending" @@ -1121,7 +1242,7 @@ mod tests { provider_instance_id: "provider-a".to_string(), display_name: "Shared files".to_string(), ttl_seconds: Some(600), - permission: server_api::ExternalWorkdirPermission::ReadOnly, + permissions: server_api::ExternalWorkdirPermissions::default(), }, ) .await diff --git a/docs/development/external-workdir-sharing.md b/docs/development/external-workdir-sharing.md new file mode 100644 index 000000000..dbe980168 --- /dev/null +++ b/docs/development/external-workdir-sharing.md @@ -0,0 +1,42 @@ +# External Workdir sharing permissions + +`yoi workdir share` exposes a client-hosted directory through an authenticated, +outbound Backend connection. The operator grants three independent categories: + +- **READ**: file reads and read-only discovery operations such as Stat, List, + ViewImage, Glob, and Grep. +- **WRITE**: Write and Edit. WRITE requires READ because updating an existing + file uses the read-before-write content-hash contract. Selecting WRITE never + enables COMMAND. +- **COMMAND**: command start, status, bounded output, and cancellation. COMMAND + does not enable READ or WRITE. It executes on the sharing host with the OS + authority of the user running the CLI; External Workdir command execution is + **not a sandbox** and does not confine the process to the shared directory. + +READ is enabled by default and COMMAND is disabled by default. Interactive mode +asks for a comma-separated category list and shows the exact grant before it is +created. For automation, repeat `--permission`: + +```sh +yoi --backend https://example.invalid workdir share ./project \ + --workspace-id workspace-a \ + --permission read \ + --permission write \ + --permission command \ + --non-interactive +``` + +`--read-only` remains an alias for READ. `--read-write` remains an alias for +READ+WRITE and never enables COMMAND. These legacy flags conflict with explicit +`--permission` arguments so ambiguous automation fails closed. + +Omitting `--ttl` creates no automatic expiry. An explicit TTL must be between +60 seconds and 24 hours. Ctrl-C, revoke, expiry, provider disconnect, attachment +detach, and provider-generation replacement terminate and reap active commands. +A reconnect creates a fresh command namespace; old handles and results are not +accepted by the new generation. + +Command output uses the existing bounded inline contract. Worker-local spill +paths are not sent to the sharing provider, and provider-local paths are never +returned to the Backend or model. Truncated External command output is therefore +reported without a full-output host path. diff --git a/openapi/server-api.json b/openapi/server-api.json index 78d8960f5..f04307fe8 100644 --- a/openapi/server-api.json +++ b/openapi/server-api.json @@ -2015,8 +2015,8 @@ "display_name": { "type": "string" }, - "permission": { - "$ref": "#/components/schemas/ExternalWorkdirPermission" + "permissions": { + "$ref": "#/components/schemas/ExternalWorkdirPermissions" }, "provider_instance_id": { "type": "string" @@ -2034,7 +2034,7 @@ "required": [ "provider_instance_id", "display_name", - "permission" + "permissions" ], "type": "object" }, @@ -2060,8 +2060,8 @@ "grant_id": { "type": "string" }, - "permission": { - "$ref": "#/components/schemas/ExternalWorkdirPermission" + "permissions": { + "$ref": "#/components/schemas/ExternalWorkdirPermissions" }, "provider_instance_id": { "type": "string" @@ -2082,18 +2082,34 @@ "working_directory_id", "provider_instance_id", "display_name", - "permission", + "permissions", "generation", "status" ], "type": "object" }, - "ExternalWorkdirPermission": { - "enum": [ - "read_only", - "read_write" + "ExternalWorkdirPermissions": { + "additionalProperties": false, + "properties": { + "command": { + "description": "Execute commands as the sharing CLI user on the provider host. This is\nexplicit authority and does not imply either filesystem category.", + "type": "boolean" + }, + "read": { + "description": "Read files and use read-only discovery operations such as Glob and Grep.", + "type": "boolean" + }, + "write": { + "description": "Create or update files with Write and Edit. Existing-file mutation also\nrequires `read` so callers can satisfy the read-before-write contract.", + "type": "boolean" + } + }, + "required": [ + "read", + "write", + "command" ], - "type": "string" + "type": "object" }, "ExtractAudit": { "properties": { @@ -13994,7 +14010,7 @@ "info": { "title": "Yoi Server API", "version": "0.1.0", - "x-yoi-source-digest": "sha256:845d8859bab4de95f0539969ae8c84bbf3fb4a0548638e4278d6c90f9bb6cb3a" + "x-yoi-source-digest": "sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be" }, "openapi": "3.1.0", "paths": { diff --git a/web/workspace/src/lib/generated/companion-api.ts b/web/workspace/src/lib/generated/companion-api.ts index def8cb001..2020f3a12 100644 --- a/web/workspace/src/lib/generated/companion-api.ts +++ b/web/workspace/src/lib/generated/companion-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-companion-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:845d8859bab4de95f0539969ae8c84bbf3fb4a0548638e4278d6c90f9bb6cb3a) +// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) // Options: companion-operation-schema-closure,openapi-3.1,strict-integer-ranges,sorted-output // Output: web/workspace/src/lib/generated/companion-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_companion_api_types -- --check diff --git a/web/workspace/src/lib/generated/repository-api.ts b/web/workspace/src/lib/generated/repository-api.ts index f7baeefdf..cb5c39c72 100644 --- a/web/workspace/src/lib/generated/repository-api.ts +++ b/web/workspace/src/lib/generated/repository-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-repository-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:845d8859bab4de95f0539969ae8c84bbf3fb4a0548638e4278d6c90f9bb6cb3a) +// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) // Options: repository-operation-schema-closure,openapi-3.1,format:uint64=safe-number,strict-safe-integers,sorted-output // Output: web/workspace/src/lib/generated/repository-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_repository_openapi_types -- --check diff --git a/web/workspace/src/lib/generated/runtime-api.ts b/web/workspace/src/lib/generated/runtime-api.ts index fe0ab354c..65d6017f2 100644 --- a/web/workspace/src/lib/generated/runtime-api.ts +++ b/web/workspace/src/lib/generated/runtime-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-runtime-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:845d8859bab4de95f0539969ae8c84bbf3fb4a0548638e4278d6c90f9bb6cb3a) +// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) // Options: runtime-operation-schema-closure,openapi-3.1,strict-integer-ranges,sorted-output // Output: web/workspace/src/lib/generated/runtime-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_runtime_api_types -- --check diff --git a/web/workspace/src/lib/generated/ticket-api.ts b/web/workspace/src/lib/generated/ticket-api.ts index b45fe6e0a..2747cd40e 100644 --- a/web/workspace/src/lib/generated/ticket-api.ts +++ b/web/workspace/src/lib/generated/ticket-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-ticket-browser-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:845d8859bab4de95f0539969ae8c84bbf3fb4a0548638e4278d6c90f9bb6cb3a) +// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) // Options: ticket-objective-merge-request-browser-operation-schema-closure,openapi-3.1,strict-integer-ranges,sorted-output // Output: web/workspace/src/lib/generated/ticket-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_ticket_api_types -- --check diff --git a/web/workspace/src/lib/generated/workdir-api.ts b/web/workspace/src/lib/generated/workdir-api.ts index feaf762c5..ed5c7392e 100644 --- a/web/workspace/src/lib/generated/workdir-api.ts +++ b/web/workspace/src/lib/generated/workdir-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-workdir-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:845d8859bab4de95f0539969ae8c84bbf3fb4a0548638e4278d6c90f9bb6cb3a) +// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) // Options: workdir-operation-schema-closure,openapi-3.1,strict-integer-ranges,sorted-output // Output: web/workspace/src/lib/generated/workdir-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_workdir_api_types -- --check @@ -9,11 +9,11 @@ export type Diagnostic = { code: string; message: string; severity: DiagnosticSe export type DiagnosticSeverity = "info" | "warning" | "error"; -export type ExternalWorkdirGrantCreateRequest = { display_name: string; permission: ExternalWorkdirPermission; provider_instance_id: string; ttl_seconds?: number | null; }; +export type ExternalWorkdirGrantCreateRequest = { display_name: string; permissions: ExternalWorkdirPermissions; provider_instance_id: string; ttl_seconds?: number | null; }; -export type ExternalWorkdirGrantResponse = { display_name: string; expires_at?: string | null; generation: number; grant_id: string; permission: ExternalWorkdirPermission; provider_instance_id: string; status: string; working_directory_id: string; workspace_id: string; }; +export type ExternalWorkdirGrantResponse = { display_name: string; expires_at?: string | null; generation: number; grant_id: string; permissions: ExternalWorkdirPermissions; provider_instance_id: string; status: string; working_directory_id: string; workspace_id: string; }; -export type ExternalWorkdirPermission = "read_only" | "read_write"; +export type ExternalWorkdirPermissions = { command: boolean; read: boolean; write: boolean; }; export type RepositoryApiError = { diagnostics?: Array; error: string; message: string; }; diff --git a/web/workspace/src/lib/generated/worker-launch-api.ts b/web/workspace/src/lib/generated/worker-launch-api.ts index 59423f838..2e2c07ed4 100644 --- a/web/workspace/src/lib/generated/worker-launch-api.ts +++ b/web/workspace/src/lib/generated/worker-launch-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-worker-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:845d8859bab4de95f0539969ae8c84bbf3fb4a0548638e4278d6c90f9bb6cb3a) +// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) // Options: worker-operation-schema-closure,openapi-3.1,strict-integer-ranges,sorted-output // Output: web/workspace/src/lib/generated/worker-launch-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_worker_launch_api_types -- --check From dcf14ef565161d54a07a8d1b1cf35a373e105d44 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 12:00:36 +0900 Subject: [PATCH 03/11] fix: harden early tool call dispatch --- crates/agen/src/engine.rs | 537 ++++++++++++++---- crates/agen/src/interceptor.rs | 14 +- .../scheme/openai_responses/events.rs | 528 +++++++++++++++-- crates/agen/src/llm_client/types.rs | 30 +- crates/agen/src/timeline/mod.rs | 1 + .../agen/src/timeline/tool_call_collector.rs | 48 +- crates/agen/tests/common/mod.rs | 13 +- crates/agen/tests/parallel_execution_test.rs | 422 +++++++++++++- crates/session-store/src/logged_item.rs | 36 +- crates/session-store/src/public_snapshot.rs | 1 + crates/worker/src/ipc/interceptor.rs | 4 + crates/worker/src/spawn/registry.rs | 3 + docs/manifest.toml | 5 +- 13 files changed, 1481 insertions(+), 161 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 2ee42feaa..fa4ab25a6 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -24,7 +24,7 @@ use crate::{ ToolCallInfo, ToolResultInfo, TurnEndAction, }, llm_client::{ - ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream, + ClientError, ConfigWarning, ItemStatus, LlmClient, Request, RequestConfig, ResponseStream, ToolDefinition, error::is_retryable, event::Event, retry::RetryPolicy, transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT, }, @@ -110,33 +110,56 @@ fn materialize_provider_history(items: &[Item]) -> ProviderHistoryProjection { let mut materialized: Vec<_> = items.iter().cloned().enumerate().collect(); let mut call_order = HashMap::::new(); let mut next_call_order = 0usize; - let mut index = 0usize; - while index < materialized.len() { - match &materialized[index].1 { - Item::ToolCall { call_id, .. } => { - call_order.insert(call_id.clone(), next_call_order); - next_call_order += 1; - index += 1; - } - Item::ToolResult { .. } => { - let start = index; - while index < materialized.len() - && matches!(materialized[index].1, Item::ToolResult { .. }) - { - index += 1; - } - materialized[start..index].sort_by_key(|(_, item)| match item { - Item::ToolResult { call_id, .. } => { - call_order.get(call_id).copied().unwrap_or(usize::MAX) - } - _ => unreachable!("tool-result run contains only ToolResult items"), - }); - } - _ => index += 1, + for (_, item) in &materialized { + if let Item::ToolCall { + call_id, + call_index, + .. + } = item + { + call_order.insert(call_id.clone(), call_index.unwrap_or(next_call_order)); + next_call_order += 1; } } + let mut index = 0usize; + while index < materialized.len() { + if !matches!( + materialized[index].1, + Item::ToolCall { .. } | Item::ToolResult { .. } + ) { + index += 1; + continue; + } + let start = index; + while index < materialized.len() + && matches!( + materialized[index].1, + Item::ToolCall { .. } | Item::ToolResult { .. } + ) + { + index += 1; + } + materialized[start..index].sort_by_key(|(_, item)| match item { + Item::ToolCall { + call_id, + call_index, + .. + } => ( + 0usize, + call_index + .or_else(|| call_order.get(call_id).copied()) + .unwrap_or(usize::MAX), + ), + Item::ToolResult { call_id, .. } => ( + 1usize, + call_order.get(call_id).copied().unwrap_or(usize::MAX), + ), + _ => unreachable!("tool run contains only calls and results"), + }); + } + let mut original_to_projected_index = vec![0; materialized.len()]; for (projected_index, (original_index, _)) in materialized.iter().enumerate() { original_to_projected_index[*original_index] = projected_index; @@ -230,8 +253,7 @@ type ToolCallExecutionInfo = ( /// provider-confirmed blocks complete. struct EarlyToolExecutionBatch { batch_id: String, - next_call_index: usize, - calls: Vec, + calls: Vec<(usize, ToolCall)>, seen_call_ids: HashSet, call_info_map: HashMap, started_calls: Vec<(String, String)>, @@ -247,7 +269,6 @@ impl EarlyToolExecutionBatch { fn new(batch_id: String) -> Self { Self { batch_id, - next_call_index: 0, calls: Vec::new(), seen_call_ids: HashSet::new(), call_info_map: HashMap::new(), @@ -378,6 +399,37 @@ struct StreamResponseOutput { early_tools: Option, } +const RESPONSE_STREAM_PUMP_CAPACITY: usize = 64; + +struct ResponseStreamPump { + receiver: mpsc::Receiver>, + task: tokio::task::JoinHandle<()>, +} + +impl ResponseStreamPump { + fn start(mut stream: ResponseStream) -> Self { + let (sender, receiver) = mpsc::channel(RESPONSE_STREAM_PUMP_CAPACITY); + let task = tokio::spawn(async move { + while let Some(event) = stream.next().await { + if sender.send(event).await.is_err() { + break; + } + } + }); + Self { receiver, task } + } + + async fn recv(&mut self) -> Option> { + self.receiver.recv().await + } +} + +impl Drop for ResponseStreamPump { + fn drop(&mut self) { + self.task.abort(); + } +} + pub struct Engine { /// LLM client client: C, @@ -1229,67 +1281,88 @@ impl Engine { exit } - /// Check for pending tool calls (for resuming from Pause) - fn get_pending_tool_calls(&self, history: &History) -> Option> { - // Find the last ToolCall items that don't have corresponding ToolResult + /// Classify unanswered tool calls for safe resume. Calls with a persisted + /// execution identity were already handed to a provider and are never rerun. + fn get_unanswered_tool_calls(&self, history: &History) -> (Vec, Vec) { let mut pending_calls = Vec::new(); + let mut started_call_ids = Vec::new(); let mut answered_call_ids = std::collections::HashSet::new(); - // First pass: collect all answered call IDs for item in history.items() { if let Item::ToolResult { call_id, .. } = item { answered_call_ids.insert(call_id.clone()); } } - // Second pass: find unanswered tool calls for item in history.items() { if let Item::ToolCall { call_id, name, arguments, + execution_id, + status, .. } = item + && !answered_call_ids.contains(call_id) { - if !answered_call_ids.contains(call_id) { - let input = serde_json::from_str::(arguments) - .ok() - .filter(Value::is_object) - .unwrap_or(Value::Null); - pending_calls.push(ToolCall { - id: call_id.clone(), - name: name.clone(), - input, - }); + if execution_id.is_some() || *status == Some(ItemStatus::InProgress) { + started_call_ids.push(call_id.clone()); + continue; } + let input = serde_json::from_str::(arguments) + .ok() + .filter(Value::is_object) + .unwrap_or(Value::Null); + pending_calls.push(ToolCall { + id: call_id.clone(), + name: name.clone(), + input, + }); } } - if pending_calls.is_empty() { - None + (pending_calls, started_call_ids) + } + + fn append_early_tool_call( + &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + call: &ToolCall, + call_index: usize, + execution_id: Option, + ) -> Result<(), EngineError> { + let status = if execution_id.is_some() { + ItemStatus::InProgress } else { - Some(pending_calls) - } + ItemStatus::Completed + }; + let item = Item::tool_call_json(&call.id, &call.name, call.input.clone()) + .with_tool_execution_metadata(call_index, execution_id) + .with_status(status); + self.append_history_items(history, std::iter::once(item), annotate) } async fn admit_early_tool_calls( &mut self, - history: &History, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, batch: &mut EarlyToolExecutionBatch, - tool_calls: Vec, + tool_calls: Vec, ) -> Result { - for mut tool_call in tool_calls { + for collected in tool_calls { + let call_index = collected.call_index; + let mut tool_call = collected.call; if !batch.seen_call_ids.insert(tool_call.id.clone()) { warn!(call_id = %tool_call.id, "Ignoring duplicate completed tool-call event"); continue; } - let call_index = batch.next_call_index; - batch.next_call_index += 1; - batch.calls.push(tool_call.clone()); + batch.calls.push((call_index, tool_call.clone())); let expected_tool_use_id = tool_call.id.clone(); let context = ToolExecutionContext::new(&tool_call.id, &batch.batch_id, call_index); if !tool_call.input.is_object() { + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; batch.synthetic_results.push_back(ToolResult::error( &tool_call.id, "Tool arguments must be one complete JSON object", @@ -1299,21 +1372,11 @@ impl Engine { } let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) else { - let attempt_id = context.batch_id.clone(); - batch - .attempt_fence - .register(tool_call.id.clone(), attempt_id.clone()); - batch - .started_calls - .push((tool_call.id.clone(), attempt_id.clone())); - let call_id = tool_call.id.clone(); - let name = tool_call.name.clone(); - batch.futures.push(Box::pin(async move { - ( - attempt_id, - ToolResult::error(&call_id, format!("Tool not found: {name}")), - ) - })); + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + batch.synthetic_results.push_back(ToolResult::error( + &tool_call.id, + format!("Tool not found: {}", tool_call.name), + )); batch.dispatched_any = true; continue; }; @@ -1332,16 +1395,28 @@ impl Engine { tool, context, }; - let action = self - .interceptor - .pre_tool_call(&mut info) - .await - .map_err(|error| { - EngineError::from(InterceptorFailure::new( - InterceptorPhase::PreToolCall, - error, - )) - })?; + let action = tokio::select! { + result = self.interceptor.pre_tool_call(&mut info) => { + result.map_err(|error| { + EngineError::from(InterceptorFailure::new( + InterceptorPhase::PreToolCall, + error, + )) + })? + } + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!(call_id = %expected_tool_use_id, "Paused while awaiting pre-tool policy"); + } + return Ok(EarlyToolAdmission::Pause); + } + cancel = self.cancel_rx.recv() => { + if cancel.is_some() { + info!(call_id = %expected_tool_use_id, "Cancelled while awaiting pre-tool policy"); + } + return Err(EngineError::Cancelled); + } + }; if info.call.id != expected_tool_use_id { return Err(InterceptorFailure::new( InterceptorPhase::PreToolCall, @@ -1352,22 +1427,30 @@ impl Engine { ) .into()); } - tool_call = info.call; - let mut context = info.context; + let ToolCallInfo { + call, + meta, + tool, + mut context, + .. + } = info; + tool_call = call; context.call_id = tool_call.id.clone(); batch.call_info_map.insert( tool_call.id.clone(), - ( - tool_call.clone(), - info.meta.clone(), - info.tool.clone(), - context.clone(), - ), + (tool_call.clone(), meta, tool.clone(), context.clone()), ); match action { PreToolAction::Continue => { - let attempt_id = context.batch_id.clone(); + let attempt_id = context.execution_id(); + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + Some(attempt_id.clone()), + )?; batch .attempt_fence .register(tool_call.id.clone(), attempt_id.clone()); @@ -1376,8 +1459,7 @@ impl Engine { .push((tool_call.id.clone(), attempt_id.clone())); let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default(); let call_id = tool_call.id.clone(); - let (handle, terminal) = - ToolExecutionHandle::start(info.tool, input_json, context); + let (handle, terminal) = ToolExecutionHandle::start(tool, input_json, context); batch.execution_handles.insert(call_id.clone(), handle); batch.futures.push(Box::pin(async move { let result = match terminal.await { @@ -1420,17 +1502,57 @@ impl Engine { ) .into()); } + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; batch.synthetic_results.push_back(result); batch.dispatched_any = true; } - PreToolAction::Skip => {} - PreToolAction::Abort(reason) => return Err(EngineError::Aborted(reason)), - PreToolAction::Pause => return Ok(EarlyToolAdmission::Pause), + PreToolAction::Skip => { + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + batch.synthetic_results.push_back(ToolResult::error( + &tool_call.id, + "Tool execution skipped by interceptor", + )); + batch.dispatched_any = true; + } + PreToolAction::Abort(reason) => { + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + batch.synthetic_results.push_back(ToolResult::error( + &tool_call.id, + format!("Tool execution aborted: {reason}"), + )); + batch.dispatched_any = true; + return Err(EngineError::Aborted(reason)); + } + PreToolAction::Pause => { + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + return Ok(EarlyToolAdmission::Pause); + } } } Ok(EarlyToolAdmission::Continue) } + async fn commit_early_synthetic_results( + &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + batch: &mut EarlyToolExecutionBatch, + ) -> Result<(), EngineError> { + while let Some(result) = batch.synthetic_results.pop_front() { + self.finalize_and_commit_tool_result( + history, + annotate, + result, + None, + &batch.call_info_map, + &mut batch.attempt_fence, + &mut batch.terminal_call_ids, + ) + .await?; + } + Ok(()) + } + async fn cancel_and_terminalize_early_batch( &mut self, history: &mut History, @@ -1439,6 +1561,25 @@ impl Engine { ) -> Option { use futures::StreamExt as _; + let mut batch_error = None; + while let Some(result) = batch.synthetic_results.pop_front() { + if let Err(error) = self + .finalize_and_commit_tool_result( + history, + annotate, + result, + None, + &batch.call_info_map, + &mut batch.attempt_fence, + &mut batch.terminal_call_ids, + ) + .await + && batch_error.is_none() + { + batch_error = Some(error); + } + } + let cancellation_request_deadline = TokioInstant::now() + self.tool_execution_policy.cancellation_request_timeout; let cancellation_requests = batch @@ -1462,7 +1603,6 @@ impl Engine { } } - let mut batch_error = None; let deadline = TokioInstant::now() + self.tool_execution_policy.terminal_confirmation_timeout; while !batch.futures.is_empty() { @@ -2131,11 +2271,33 @@ impl Engine { "Starting engine run" ); - // Resume pending tool calls from a previous Pause - if let Some(tool_calls) = self.get_pending_tool_calls(history) { - info!("Resuming pending tool calls"); + // Resume only calls durably proven not started. A persisted started + // attempt without a terminal result is closed as OutcomeUnknown first. + let (pending_tool_calls, started_call_ids) = self.get_unanswered_tool_calls(history); + if !started_call_ids.is_empty() { + let call_info_map = HashMap::new(); + let mut attempt_fence = ToolExecutionAttemptFence::default(); + let mut terminal_call_ids = HashSet::new(); + for call_id in started_call_ids { + self.finalize_and_commit_tool_result( + history, + annotate, + ToolResult::outcome_unknown(&call_id), + None, + &call_info_map, + &mut attempt_fence, + &mut terminal_call_ids, + ) + .await?; + } + } + if !pending_tool_calls.is_empty() { + info!( + count = pending_tool_calls.len(), + "Resuming pending tool calls" + ); if let Some(result) = self - .execute_and_commit_tools(history, annotate, tool_calls) + .execute_and_commit_tools(history, annotate, pending_tool_calls) .await? { return Ok(result); @@ -2348,17 +2510,28 @@ impl Engine { self.request_trace_payload(&request), ); let request = self.attach_transport_trace(request, current_turn, current_llm_call); + let provider_supports_early = self.client.tool_call_completion_support() + == crate::llm_client::ToolCallCompletionSupport::PerBlock; + let interceptor_supports_early = self.interceptor.supports_early_tool_dispatch(); let early_dispatch = self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete - && self.client.tool_call_completion_support() - == crate::llm_client::ToolCallCompletionSupport::PerBlock; + && provider_supports_early + && interceptor_supports_early; if self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete - && !early_dispatch + && !provider_supports_early { self.emit_warning( "tool-call early dispatch deferred until response completion: provider does not expose a trustworthy per-call completion boundary", ); } + if self.tool_call_dispatch_mode == ToolCallDispatchMode::OnToolCallComplete + && provider_supports_early + && !interceptor_supports_early + { + self.emit_warning( + "tool-call early dispatch deferred until response completion: interceptor requires the whole-response turn-end gate", + ); + } let assistant_start = history.len(); let StreamResponseOutput { completion: stream_outcome, @@ -2466,9 +2639,25 @@ impl Engine { // `append_history_items` so observers see each item as it lands. let reasoning_items = self.thinking_block_collector.take_collected(); let text_blocks = self.text_block_collector.take_collected(); - let remaining_tool_calls = self.tool_call_collector.take_collected(); - let assistant_items = - self.build_assistant_items(&reasoning_items, &text_blocks, &remaining_tool_calls); + let remaining_collected_calls = self.tool_call_collector.take_collected_with_index(); + let mut remaining_tool_calls = remaining_collected_calls + .iter() + .map(|collected| (collected.call_index, collected.call.clone())) + .collect::>(); + remaining_tool_calls.sort_by_key(|(call_index, _)| *call_index); + let remaining_tool_calls = remaining_tool_calls + .into_iter() + .map(|(_, call)| call) + .collect::>(); + let assistant_items = self.build_assistant_items( + &reasoning_items, + &text_blocks, + if early_dispatch { + &[] + } else { + &remaining_tool_calls + }, + ); if let Err(error) = self.append_history_items(history, assistant_items, annotate) { if let Some(batch) = early_tools.take() && let Err(cleanup_error) = self @@ -2489,8 +2678,9 @@ impl Engine { let admission = self .admit_early_tool_calls( history, + annotate, early_tools.as_mut().expect("batch initialized"), - remaining_tool_calls.clone(), + remaining_collected_calls, ) .await; let admission = match admission { @@ -2513,10 +2703,27 @@ impl Engine { } return Ok(EngineResult::Paused); } + if let Some(batch) = early_tools.as_mut() + && let Err(error) = self + .commit_early_synthetic_results(history, annotate, batch) + .await + { + if let Some(mut batch) = early_tools.take() { + let cleanup_error = self + .cancel_and_terminalize_early_batch(history, annotate, &mut batch) + .await; + return Err(cleanup_error.unwrap_or(error)); + } + return Err(error); + } } let mut tool_calls = early_tools .as_ref() - .map(|batch| batch.calls.clone()) + .map(|batch| { + let mut calls = batch.calls.clone(); + calls.sort_by_key(|(call_index, _)| *call_index); + calls.into_iter().map(|(_, call)| call).collect() + }) .unwrap_or_default(); if !early_dispatch { tool_calls = remaining_tool_calls; @@ -2797,13 +3004,15 @@ impl Engine { "Sending request to LLM" ); - let mut stream = self.open_stream_with_retry(request, turn, llm_call).await?; + let stream = self.open_stream_with_retry(request, turn, llm_call).await?; + self.tool_call_collector.begin_response(); + let mut stream = ResponseStreamPump::start(stream); let mut early_tools: Option = None; let mut event_count: usize = 0; loop { tokio::select! { - event_result = stream.next() => { + event_result = stream.recv() => { match event_result { Some(result) => { match &result { @@ -2839,7 +3048,7 @@ impl Engine { self.emit_stream_event(turn, llm_call, &event); self.timeline.dispatch(&event); if early_dispatch { - let collected = self.tool_call_collector.take_collected(); + let collected = self.tool_call_collector.take_collected_with_index(); if !collected.is_empty() { if early_tools.is_none() { let batch_id = format!( @@ -2854,22 +3063,23 @@ impl Engine { let batch = early_tools.as_ref().expect("batch initialized"); collected .into_iter() - .filter(|call| { - !batch.seen_call_ids.contains(&call.id) + .filter(|collected| { + !batch.seen_call_ids.contains(&collected.call.id) }) .collect::>() }; if !fresh_calls.is_empty() { - // Commit every completed assistant block observed so far, - // including the ToolCall itself, before provider ownership - // can start. This is the durable fail-closed gate. + // Commit every completed reasoning/text block observed so + // far. Admission durably commits each ToolCall with its + // started/not-started state immediately before ownership + // can transfer to a provider. let reasoning = self.thinking_block_collector.take_collected(); let text = self.text_block_collector.take_collected(); let assistant_items = self.build_assistant_items( &reasoning, &text, - &fresh_calls, + &[], ); if let Err(error) = self.append_history_items( history, @@ -2892,13 +3102,38 @@ impl Engine { early_tools.as_mut().expect("batch initialized"); self.admit_early_tool_calls( history, + annotate, batch, fresh_calls, ) .await }; match admission { - Ok(EarlyToolAdmission::Continue) => {} + Ok(EarlyToolAdmission::Continue) => { + let batch = early_tools + .as_mut() + .expect("batch initialized"); + if let Err(error) = self + .commit_early_synthetic_results( + history, + annotate, + batch, + ) + .await + { + if let Some(mut batch) = early_tools.take() { + let cleanup_error = self + .cancel_and_terminalize_early_batch( + history, + annotate, + &mut batch, + ) + .await; + return Err(cleanup_error.unwrap_or(error)); + } + return Err(error); + } + } Ok(EarlyToolAdmission::Pause) => { self.timeline.abort_current_block(); self.timeline.flush_usage(); @@ -2959,6 +3194,45 @@ impl Engine { None => break, } } + terminal = async { + early_tools + .as_mut() + .expect("terminal branch requires early batch") + .futures + .next() + .await + }, if early_tools.as_ref().is_some_and(|batch| !batch.futures.is_empty()) => { + let (attempt_id, result) = + terminal.expect("non-empty early terminal queue returns a result"); + let commit = { + let batch = early_tools.as_mut().expect("batch initialized"); + self.finalize_and_commit_tool_result( + history, + annotate, + result, + Some(&attempt_id), + &batch.call_info_map, + &mut batch.attempt_fence, + &mut batch.terminal_call_ids, + ) + .await + }; + if let Err(error) = commit { + self.timeline.abort_current_block(); + self.timeline.flush_usage(); + if let Some(mut batch) = early_tools.take() { + let cleanup_error = self + .cancel_and_terminalize_early_batch( + history, + annotate, + &mut batch, + ) + .await; + return Err(cleanup_error.unwrap_or(error)); + } + return Err(error); + } + } pause = self.pause_rx.recv() => { if pause.is_some() { info!("Paused during response stream"); @@ -3716,6 +3990,41 @@ mod tests { assert_eq!(projection.original_to_projected_index, [0, 1, 3, 2]); } + #[test] + fn provider_projection_groups_interleaved_calls_and_results_in_model_order() { + let items = vec![ + Item::tool_call_json("call_second", "second", serde_json::json!({})) + .with_tool_execution_metadata(1, Some("batch:call_second".to_string())) + .with_status(ItemStatus::InProgress), + Item::tool_result("call_second", "second result"), + Item::tool_call_json("call_first", "first", serde_json::json!({})) + .with_tool_execution_metadata(0, Some("batch:call_first".to_string())) + .with_status(ItemStatus::InProgress), + Item::tool_result("call_first", "first result"), + ]; + + let projection = materialize_provider_history(&items); + let order = projection + .items + .iter() + .map(|item| match item { + Item::ToolCall { call_id, .. } => format!("call:{call_id}"), + Item::ToolResult { call_id, .. } => format!("result:{call_id}"), + _ => unreachable!(), + }) + .collect::>(); + assert_eq!( + order, + [ + "call:call_first", + "call:call_second", + "result:call_first", + "result:call_second", + ] + ); + assert_eq!(projection.original_to_projected_index, [1, 3, 0, 2]); + } + #[test] fn tool_attachment_round_trips_through_durable_history_json() { let body: Arc<[u8]> = Arc::from(&b"image-body"[..]); diff --git a/crates/agen/src/interceptor.rs b/crates/agen/src/interceptor.rs index 28b2afa6b..67d34166d 100644 --- a/crates/agen/src/interceptor.rs +++ b/crates/agen/src/interceptor.rs @@ -336,6 +336,14 @@ pub struct ToolResultInfo<'a, A = ()> { /// flows, permission checks, and other trusted host adaptation. #[async_trait] pub trait Interceptor: Send + Sync { + /// Whether side effects may start before [`Self::on_assistant_turn_end`]. + /// + /// The conservative default preserves the historical whole-response gate. + /// Interceptors that do not use turn-end policy may opt in explicitly. + fn supports_early_tool_dispatch(&self) -> bool { + false + } + /// Called after receiving user input, before adding it to Engine history. async fn on_prompt_submit( &self, @@ -420,4 +428,8 @@ pub trait Interceptor: Send + Sync { pub(crate) struct DefaultInterceptor; #[async_trait] -impl Interceptor for DefaultInterceptor {} +impl Interceptor for DefaultInterceptor { + fn supports_early_tool_dispatch(&self) -> bool { + true + } +} diff --git a/crates/agen/src/llm_client/scheme/openai_responses/events.rs b/crates/agen/src/llm_client/scheme/openai_responses/events.rs index 47c8233b4..cb6ead57e 100644 --- a/crates/agen/src/llm_client/scheme/openai_responses/events.rs +++ b/crates/agen/src/llm_client/scheme/openai_responses/events.rs @@ -28,6 +28,43 @@ pub struct OpenAIResponsesState { /// 蓄積、`output_item.done`(Reasoning) で既存 reasoning_text block または /// metadata-only Thinking block に reasoning persistence material を載せる。 pending_reasoning: HashMap, + /// Tool input deltas and authoritative final values, keyed by output index. + pending_tool_inputs: HashMap, +} + +/// A tool input stream accumulated until an authoritative done payload arrives. +#[derive(Debug)] +struct PendingToolInput { + kind: ToolInputKind, + streamed: String, + saw_delta: bool, + final_value: Option, +} + +impl PendingToolInput { + fn new(kind: ToolInputKind) -> Self { + Self { + kind, + streamed: String::new(), + saw_delta: false, + final_value: None, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ToolInputKind { + FunctionCall, + CustomToolCall, +} + +impl ToolInputKind { + fn label(self) -> &'static str { + match self { + Self::FunctionCall => "function-call arguments", + Self::CustomToolCall => "custom-tool input", + } + } } /// 1 つの reasoning output_item の蓄積バッファ。 @@ -173,6 +210,8 @@ enum OutputItem { id: Option, call_id: String, name: String, + #[serde(default)] + arguments: Option, }, CustomToolCall { #[allow(dead_code)] @@ -180,6 +219,8 @@ enum OutputItem { id: Option, call_id: String, name: String, + #[serde(default)] + input: Option, }, #[serde(other)] Other, @@ -260,12 +301,24 @@ struct FunctionCallArgumentsDelta { delta: String, } +#[derive(Debug, Deserialize)] +struct FunctionCallArgumentsDone { + output_index: usize, + arguments: String, +} + #[derive(Debug, Deserialize)] struct CustomToolCallInputDelta { output_index: usize, delta: String, } +#[derive(Debug, Deserialize)] +struct CustomToolCallInputDone { + output_index: usize, + input: String, +} + #[derive(Debug, Deserialize)] struct ResponseCompleted { response: CompletedResponse, @@ -398,16 +451,20 @@ pub(crate) fn parse_sse( "response.output_item.added" => { let ev: OutputItemAdded = from_json(data)?; match ev.item { - OutputItem::FunctionCall { call_id, name, .. } - | OutputItem::CustomToolCall { call_id, name, .. } => { - let info = - state.allocate(SlotKey::OutputItem(ev.output_index), BlockType::ToolUse); - Ok(vec![Event::BlockStart(BlockStart { - index: info.flat_index, - block_type: BlockType::ToolUse, - metadata: BlockMetadata::ToolUse { id: call_id, name }, - })]) - } + OutputItem::FunctionCall { call_id, name, .. } => start_tool_output_item( + state, + ev.output_index, + ToolInputKind::FunctionCall, + call_id, + name, + ), + OutputItem::CustomToolCall { call_id, name, .. } => start_tool_output_item( + state, + ev.output_index, + ToolInputKind::CustomToolCall, + call_id, + name, + ), OutputItem::Reasoning { id, .. } => { // wrapper を確保。中身の content_part / summary_part は // 別 SlotKey で扱われ続ける（Streaming 表示は維持）。 @@ -431,19 +488,19 @@ pub(crate) fn parse_sse( id, encrypted_content, .. - } = ev.item + } = &ev.item { let mut pending = state .pending_reasoning .remove(&ev.output_index) .unwrap_or_default(); if pending.id.is_none() { - pending.id = id; + pending.id = id.clone(); } let mut stop_blocks = std::mem::take(&mut pending.deferred_thinking_stops); stop_blocks.extend(state.take_active_reasoning_slots(ev.output_index)); - let reasoning = pending.into_reasoning_data(encrypted_content); + let reasoning = pending.into_reasoning_data(encrypted_content.clone()); if stop_blocks.is_empty() { let info = @@ -478,15 +535,45 @@ pub(crate) fn parse_sse( }) .collect()); } - if let Some(info) = state.slots.remove(&SlotKey::OutputItem(ev.output_index)) { - Ok(vec![Event::BlockStop(BlockStop { - index: info.flat_index, - block_type: info.block_type, - stop_reason: None, - reasoning: None, - })]) - } else { - Ok(Vec::new()) + match ev.item { + OutputItem::FunctionCall { + call_id, + name, + arguments, + .. + } => finish_tool_output_item( + state, + ev.output_index, + ToolInputKind::FunctionCall, + call_id, + name, + arguments, + ), + OutputItem::CustomToolCall { + call_id, + name, + input, + .. + } => finish_tool_output_item( + state, + ev.output_index, + ToolInputKind::CustomToolCall, + call_id, + name, + input, + ), + _ => { + if let Some(info) = state.slots.remove(&SlotKey::OutputItem(ev.output_index)) { + Ok(vec![Event::BlockStop(BlockStop { + index: info.flat_index, + block_type: info.block_type, + stop_reason: None, + reasoning: None, + })]) + } else { + Ok(Vec::new()) + } + } } } @@ -619,30 +706,50 @@ pub(crate) fn parse_sse( "response.function_call_arguments.delta" => { let ev: FunctionCallArgumentsDelta = from_json(data)?; - Ok(ensure_and_delta( + record_tool_input_delta( state, - SlotKey::OutputItem(ev.output_index), - BlockType::ToolUse, + ev.output_index, + ToolInputKind::FunctionCall, + ev.delta, + ) + } + + "response.function_call_arguments.done" => { + let ev: FunctionCallArgumentsDone = from_json(data)?; + reconcile_tool_input( + state, + ev.output_index, + ToolInputKind::FunctionCall, + ev.arguments, BlockMetadata::ToolUse { id: String::new(), name: String::new(), }, - DeltaContent::InputJson(ev.delta), - )) + ) } "response.custom_tool_call_input.delta" => { let ev: CustomToolCallInputDelta = from_json(data)?; - Ok(ensure_and_delta( + record_tool_input_delta( state, - SlotKey::OutputItem(ev.output_index), - BlockType::ToolUse, + ev.output_index, + ToolInputKind::CustomToolCall, + ev.delta, + ) + } + + "response.custom_tool_call_input.done" => { + let ev: CustomToolCallInputDone = from_json(data)?; + reconcile_tool_input( + state, + ev.output_index, + ToolInputKind::CustomToolCall, + ev.input, BlockMetadata::ToolUse { id: String::new(), name: String::new(), }, - DeltaContent::InputJson(ev.delta), - )) + ) } "error" => { @@ -666,6 +773,181 @@ pub(crate) fn parse_sse( } } +fn start_tool_output_item( + state: &mut OpenAIResponsesState, + output_index: usize, + kind: ToolInputKind, + call_id: String, + name: String, +) -> Result, ClientError> { + ensure_tool_input_kind(state, output_index, kind)?; + let (info, just_created) = + state.get_or_allocate(SlotKey::OutputItem(output_index), BlockType::ToolUse); + if just_created { + Ok(vec![Event::BlockStart(BlockStart { + index: info.flat_index, + block_type: BlockType::ToolUse, + metadata: BlockMetadata::ToolUse { id: call_id, name }, + })]) + } else { + Ok(Vec::new()) + } +} + +fn record_tool_input_delta( + state: &mut OpenAIResponsesState, + output_index: usize, + kind: ToolInputKind, + delta: String, +) -> Result, ClientError> { + let pending = ensure_tool_input_kind(state, output_index, kind)?; + if pending.final_value.is_some() { + return Err(tool_input_error( + "tool_input_delta_after_done", + format!( + "OpenAI Responses {} delta arrived after its final value at output index {output_index}", + kind.label() + ), + )); + } + pending.saw_delta = true; + pending.streamed.push_str(&delta); + + Ok(ensure_and_delta( + state, + SlotKey::OutputItem(output_index), + BlockType::ToolUse, + BlockMetadata::ToolUse { + id: String::new(), + name: String::new(), + }, + DeltaContent::InputJson(delta), + )) +} + +fn reconcile_tool_input( + state: &mut OpenAIResponsesState, + output_index: usize, + kind: ToolInputKind, + final_value: String, + metadata: BlockMetadata, +) -> Result, ClientError> { + let pending = ensure_tool_input_kind(state, output_index, kind)?; + if pending + .final_value + .as_ref() + .is_some_and(|existing| existing != &final_value) + { + return Err(tool_input_error( + "tool_input_final_mismatch", + format!( + "OpenAI Responses authoritative {} values disagree at output index {output_index}", + kind.label() + ), + )); + } + if pending.saw_delta && pending.streamed != final_value { + return Err(tool_input_error( + "tool_input_delta_mismatch", + format!( + "OpenAI Responses streamed {} does not match its authoritative final value at output index {output_index}", + kind.label() + ), + )); + } + + let should_emit = !pending.saw_delta; + if should_emit { + pending.saw_delta = true; + pending.streamed.clone_from(&final_value); + } + pending.final_value = Some(final_value.clone()); + + if should_emit { + Ok(ensure_and_delta( + state, + SlotKey::OutputItem(output_index), + BlockType::ToolUse, + metadata, + DeltaContent::InputJson(final_value), + )) + } else { + Ok(Vec::new()) + } +} + +fn finish_tool_output_item( + state: &mut OpenAIResponsesState, + output_index: usize, + kind: ToolInputKind, + call_id: String, + name: String, + final_value: Option, +) -> Result, ClientError> { + let metadata = BlockMetadata::ToolUse { id: call_id, name }; + let mut out = if let Some(final_value) = final_value { + reconcile_tool_input(state, output_index, kind, final_value, metadata.clone())? + } else { + let pending = ensure_tool_input_kind(state, output_index, kind)?; + if pending.final_value.is_none() { + return Err(tool_input_error( + "tool_input_final_missing", + format!( + "OpenAI Responses {} ended without an authoritative final value at output index {output_index}", + kind.label() + ), + )); + } + Vec::new() + }; + + let (info, just_created) = + state.get_or_allocate(SlotKey::OutputItem(output_index), BlockType::ToolUse); + if just_created { + out.push(Event::BlockStart(BlockStart { + index: info.flat_index, + block_type: BlockType::ToolUse, + metadata, + })); + } + state.slots.remove(&SlotKey::OutputItem(output_index)); + state.pending_tool_inputs.remove(&output_index); + out.push(Event::BlockStop(BlockStop { + index: info.flat_index, + block_type: info.block_type, + stop_reason: None, + reasoning: None, + })); + Ok(out) +} + +fn ensure_tool_input_kind( + state: &mut OpenAIResponsesState, + output_index: usize, + kind: ToolInputKind, +) -> Result<&mut PendingToolInput, ClientError> { + let pending = state + .pending_tool_inputs + .entry(output_index) + .or_insert_with(|| PendingToolInput::new(kind)); + if pending.kind != kind { + return Err(tool_input_error( + "tool_input_kind_mismatch", + format!("OpenAI Responses tool input kind changed at output index {output_index}"), + )); + } + Ok(pending) +} + +fn tool_input_error(code: &str, message: String) -> ClientError { + ClientError::Api { + status: None, + code: Some(code.to_string()), + message, + retry_after: None, + } +} + fn response_failure_diagnostic(event_type: &str, ev: ResponseFailed) -> (Option, String) { let mut diagnostic = Map::new(); diagnostic.insert("event".to_string(), Value::String(event_type.to_string())); @@ -855,6 +1137,26 @@ mod tests { parse_sse(event_type, data, state).unwrap() } + fn input_json_deltas(events: &[Event]) -> Vec<&str> { + events + .iter() + .filter_map(|event| match event { + Event::BlockDelta(BlockDelta { + delta: DeltaContent::InputJson(delta), + .. + }) => Some(delta.as_str()), + _ => None, + }) + .collect() + } + + fn api_error_code(error: ClientError) -> Option { + match error { + ClientError::Api { code, .. } => code, + other => panic!("expected ClientError::Api, got {other:?}"), + } + } + #[test] fn created_emits_status_started() { let (events, _) = run("response.created", r#"{"response":{}}"#); @@ -981,6 +1283,19 @@ mod tests { if let Event::BlockDelta(d) = &ev[0] { assert!(matches!(&d.delta, DeltaContent::InputJson(j) if j == "{\"x\":")); } + let ev = with( + &mut state, + "response.function_call_arguments.delta", + r#"{"output_index":1,"item_id":"fc1","delta":"1}"}"#, + ); + assert_eq!(ev.len(), 1); + assert!(matches!( + &ev[0], + Event::BlockDelta(BlockDelta { + delta: DeltaContent::InputJson(j), + .. + }) if j == "1}" + )); // output_item.done → BlockStop let ev = with( &mut state, @@ -1012,6 +1327,155 @@ mod tests { } } + #[test] + fn function_call_done_payload_emits_missing_final_input_before_stop() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_1","name":"lookup"}}"#, + ); + + let events = with( + &mut state, + "response.output_item.done", + r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"city\":\"Tokyo\"}"}}"#, + ); + + assert_eq!(input_json_deltas(&events), vec![r#"{"city":"Tokyo"}"#]); + assert!(matches!(events.last(), Some(Event::BlockStop(_)))); + } + + #[test] + fn function_call_arguments_done_accepts_exact_stream_without_duplication() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_1","name":"lookup"}}"#, + ); + let delta = with( + &mut state, + "response.function_call_arguments.delta", + r#"{"output_index":0,"delta":"{\"city\":\"Tokyo\"}"}"#, + ); + assert_eq!(input_json_deltas(&delta), vec![r#"{"city":"Tokyo"}"#]); + + let done = with( + &mut state, + "response.function_call_arguments.done", + r#"{"output_index":0,"arguments":"{\"city\":\"Tokyo\"}"}"#, + ); + assert!(done.is_empty()); + let output_done = with( + &mut state, + "response.output_item.done", + r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_1","name":"lookup","arguments":"{\"city\":\"Tokyo\"}"}}"#, + ); + assert!(input_json_deltas(&output_done).is_empty()); + assert!(matches!(output_done.as_slice(), [Event::BlockStop(_)])); + } + + #[test] + fn function_call_arguments_done_rejects_stream_mismatch() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_1","name":"lookup"}}"#, + ); + with( + &mut state, + "response.function_call_arguments.delta", + r#"{"output_index":0,"delta":"{\"city\":\"Kyoto\"}"}"#, + ); + + let error = parse_sse( + "response.function_call_arguments.done", + r#"{"output_index":0,"arguments":"{\"city\":\"Tokyo\"}"}"#, + &mut state, + ) + .unwrap_err(); + assert_eq!( + api_error_code(error).as_deref(), + Some("tool_input_delta_mismatch") + ); + } + + #[test] + fn custom_tool_input_done_event_emits_missing_final_input_once() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell"}}"#, + ); + + let done = with( + &mut state, + "response.custom_tool_call_input.done", + r#"{"output_index":0,"input":"echo ready"}"#, + ); + assert_eq!(input_json_deltas(&done), vec!["echo ready"]); + let output_done = with( + &mut state, + "response.output_item.done", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell","input":"echo ready"}}"#, + ); + assert!(input_json_deltas(&output_done).is_empty()); + assert!(matches!(output_done.as_slice(), [Event::BlockStop(_)])); + } + + #[test] + fn custom_tool_done_payload_accepts_exact_stream_without_duplication() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell"}}"#, + ); + let delta = with( + &mut state, + "response.custom_tool_call_input.delta", + r#"{"output_index":0,"delta":"echo ready"}"#, + ); + assert_eq!(input_json_deltas(&delta), vec!["echo ready"]); + + let output_done = with( + &mut state, + "response.output_item.done", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell","input":"echo ready"}}"#, + ); + assert!(input_json_deltas(&output_done).is_empty()); + assert!(matches!(output_done.as_slice(), [Event::BlockStop(_)])); + } + + #[test] + fn custom_tool_done_payload_rejects_stream_mismatch() { + let mut state = OpenAIResponsesState::default(); + with( + &mut state, + "response.output_item.added", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell"}}"#, + ); + with( + &mut state, + "response.custom_tool_call_input.delta", + r#"{"output_index":0,"delta":"echo stale"}"#, + ); + + let error = parse_sse( + "response.output_item.done", + r#"{"output_index":0,"item":{"type":"custom_tool_call","call_id":"call_1","name":"shell","input":"echo ready"}}"#, + &mut state, + ) + .unwrap_err(); + assert_eq!( + api_error_code(error).as_deref(), + Some("tool_input_delta_mismatch") + ); + } + #[test] fn reasoning_text_delta_emits_thinking() { let mut state = OpenAIResponsesState::default(); diff --git a/crates/agen/src/llm_client/types.rs b/crates/agen/src/llm_client/types.rs index 101e97731..50ed7506a 100644 --- a/crates/agen/src/llm_client/types.rs +++ b/crates/agen/src/llm_client/types.rs @@ -104,7 +104,14 @@ pub enum Item { name: String, /// Tool arguments as JSON string arguments: String, - /// Item status + /// Provider/model order among tool calls in this assistant response. + #[serde(default, skip_serializing_if = "Option::is_none")] + call_index: Option, + /// Durable identity of a tool execution attempt that has started. + #[serde(default, skip_serializing_if = "Option::is_none")] + execution_id: Option, + /// Item status. `InProgress` on a persisted ToolCall means provider + /// execution ownership was transferred and must never be inferred pending. #[serde(skip_serializing_if = "Option::is_none")] status: Option, }, @@ -234,6 +241,8 @@ impl Item { call_id: call_id.into(), name: name.into(), arguments: arguments.into(), + call_index: None, + execution_id: None, status: None, } } @@ -247,6 +256,25 @@ impl Item { Self::tool_call(call_id, name, arguments.to_string()) } + /// Attach response-local ordering and, when execution has started, its + /// durable attempt identity to a tool call. + pub fn with_tool_execution_metadata( + mut self, + call_index: usize, + execution_id: Option, + ) -> Self { + if let Self::ToolCall { + call_index: item_call_index, + execution_id: item_execution_id, + .. + } = &mut self + { + *item_call_index = Some(call_index); + *item_execution_id = execution_id; + } + self + } + /// Create a tool result item with summary only (no content). pub fn tool_result(call_id: impl Into, summary: impl Into) -> Self { Self::tool_result_item(call_id, summary, None, false) diff --git a/crates/agen/src/timeline/mod.rs b/crates/agen/src/timeline/mod.rs index a1bd24cb8..b016c2592 100644 --- a/crates/agen/src/timeline/mod.rs +++ b/crates/agen/src/timeline/mod.rs @@ -21,6 +21,7 @@ pub use event::*; pub use text_block_collector::TextBlockCollector; pub use thinking_block_collector::ThinkingBlockCollector; pub use timeline::Timeline; +pub(crate) use tool_call_collector::CollectedToolCall; pub use tool_call_collector::ToolCallCollector; // 型定義からのre-export diff --git a/crates/agen/src/timeline/tool_call_collector.rs b/crates/agen/src/timeline/tool_call_collector.rs index 3c2ed0ca2..3663a8545 100644 --- a/crates/agen/src/timeline/tool_call_collector.rs +++ b/crates/agen/src/timeline/tool_call_collector.rs @@ -17,10 +17,19 @@ pub struct CollectorState { /// 現在のツール呼び出し情報 (ブロック進行中) current_id: Option, current_name: Option, + /// Zero-based order among tool calls in this assistant response. + call_index: Option, /// 蓄積中のJSON入力 input_json_buffer: String, } +/// One provider-completed call together with its stable model order. +#[derive(Debug, Clone)] +pub(crate) struct CollectedToolCall { + pub call_index: usize, + pub call: ToolCall, +} + /// ToolCallCollector - ToolUseブロックハンドラ /// /// Timelineに登録してToolUseブロックイベントを受信し、 @@ -28,7 +37,8 @@ pub struct CollectorState { #[derive(Clone)] pub struct ToolCallCollector { /// 収集されたToolCall - collected: Arc>>, + collected: Arc>>, + next_call_index: Arc>, } impl ToolCallCollector { @@ -36,18 +46,38 @@ impl ToolCallCollector { pub fn new() -> Self { Self { collected: Arc::new(Mutex::new(Vec::new())), + next_call_index: Arc::new(Mutex::new(0)), } } /// 収集されたToolCallを取得してクリア pub fn take_collected(&self) -> Vec { + self.take_collected_with_index() + .into_iter() + .map(|collected| collected.call) + .collect() + } + + /// Drain completed calls while preserving their model-returned order. + pub(crate) fn take_collected_with_index(&self) -> Vec { let mut guard = self.collected.lock().unwrap(); std::mem::take(&mut *guard) } /// 収集されたToolCallの参照を取得 pub fn collected(&self) -> Vec { - self.collected.lock().unwrap().clone() + self.collected + .lock() + .unwrap() + .iter() + .map(|collected| collected.call.clone()) + .collect() + } + + /// Start one response-local tool-call ordering scope. + pub(crate) fn begin_response(&self) { + self.collected.lock().unwrap().clear(); + *self.next_call_index.lock().unwrap() = 0; } /// 収集されたToolCallがあるかどうか @@ -75,6 +105,11 @@ impl Handler for ToolCallCollector { ToolUseBlockEvent::Start(start) => { scope.current_id = Some(start.id.clone()); scope.current_name = Some(start.name.clone()); + if scope.call_index.is_none() { + let mut next = self.next_call_index.lock().unwrap(); + scope.call_index = Some(*next); + *next += 1; + } scope.input_json_buffer.clear(); } ToolUseBlockEvent::InputJsonDelta(delta) => { @@ -98,8 +133,15 @@ impl Handler for ToolCallCollector { }; let tool_call = ToolCall { id, name, input }; + let call_index = scope + .call_index + .take() + .expect("tool-use stop follows a started collector scope"); - self.collected.lock().unwrap().push(tool_call); + self.collected.lock().unwrap().push(CollectedToolCall { + call_index, + call: tool_call, + }); } scope.input_json_buffer.clear(); } diff --git a/crates/agen/tests/common/mod.rs b/crates/agen/tests/common/mod.rs index 5f4d4e941..277f3ec13 100644 --- a/crates/agen/tests/common/mod.rs +++ b/crates/agen/tests/common/mod.rs @@ -7,7 +7,7 @@ use std::pin::Pin; use std::sync::{Arc, Mutex}; use agen::llm_client::event::{BlockType, DeltaContent, Event}; -use agen::llm_client::{ClientError, LlmClient, Request}; +use agen::llm_client::{ClientError, LlmClient, Request, ToolCallCompletionSupport}; use agen::timeline::{Handler, TextBlockEvent, TextBlockKind, Timeline}; use async_trait::async_trait; use futures::Stream; @@ -20,6 +20,7 @@ pub struct MockLlmClient { responses: Arc>>, call_count: Arc, requests: Arc>>, + completion_support: ToolCallCompletionSupport, } impl MockLlmClient { @@ -32,6 +33,7 @@ impl MockLlmClient { responses: Arc::new(responses), call_count: Arc::new(AtomicUsize::new(0)), requests: Arc::new(Mutex::new(Vec::new())), + completion_support: ToolCallCompletionSupport::ResponseComplete, } } @@ -47,6 +49,11 @@ impl MockLlmClient { pub fn requests(&self) -> Vec { self.requests.lock().unwrap().clone() } + + pub fn with_completion_support(mut self, support: ToolCallCompletionSupport) -> Self { + self.completion_support = support; + self + } } #[async_trait] @@ -55,6 +62,10 @@ impl LlmClient for MockLlmClient { Box::new(self.clone()) } + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + self.completion_support + } + async fn stream( &self, request: Request, diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index f38f9f7cc..f6d445fb0 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -7,8 +7,9 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use agen::interceptor::{ - Interceptor, InterceptorError, InterceptorErrorCategory, InterceptorPhase, InterceptorResult, - PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo, + AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorErrorCategory, + InterceptorPhase, InterceptorResult, PostToolAction, PreToolAction, ToolCallInfo, + ToolResultInfo, TurnEndAction, }; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::{ @@ -307,6 +308,46 @@ impl LlmClient for ControlledStreamClient { } } +#[derive(Clone)] +struct BoundedStreamClient { + receiver: Arc>>>>, +} + +impl BoundedStreamClient { + fn new(capacity: usize) -> (Self, tokio::sync::mpsc::Sender>) { + let (tx, rx) = tokio::sync::mpsc::channel(capacity); + ( + Self { + receiver: Arc::new(Mutex::new(Some(rx))), + }, + tx, + ) + } +} + +#[async_trait] +impl LlmClient for BoundedStreamClient { + async fn stream(&self, _request: Request) -> Result { + let mut rx = self + .receiver + .lock() + .unwrap() + .take() + .expect("bounded stream is consumed once"); + Ok(Box::pin(futures::stream::poll_fn(move |cx| { + rx.poll_recv(cx) + }))) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + ToolCallCompletionSupport::PerBlock + } +} + #[derive(Clone)] struct BarrierTool { name: &'static str, @@ -400,6 +441,141 @@ async fn default_dispatch_waits_for_response_completion() { assert_eq!(probe.starts.load(Ordering::SeqCst), 1); } +#[tokio::test] +async fn turn_end_policy_incompatibility_visibly_falls_back_before_side_effects() { + struct TurnEndPause; + + #[async_trait] + impl Interceptor for TurnEndPause { + async fn on_assistant_turn_end( + &self, + _context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { + Ok(TurnEndAction::Pause) + } + } + + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = BarrierTool::new("turn_end_fallback"); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(TurnEndPause); + engine.register_tool(tool.definition()); + let (warning_tx, mut warning_rx) = tokio::sync::mpsc::unbounded_channel(); + engine.on_warning(move |warning| { + let _ = warning_tx.send(warning.to_string()); + }); + let (ping_tx, mut ping_rx) = tokio::sync::mpsc::unbounded_channel(); + engine.on_stream_event(move |_, _, event| { + if matches!(event, Event::Ping(_)) { + let _ = ping_tx.send(()); + } + }); + + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + tx.send(Ok(Event::tool_use_start( + 0, + "call_turn_end_fallback", + "turn_end_fallback", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tx.send(Ok(Event::ping())).unwrap(); + ping_rx.recv().await.unwrap(); + assert_eq!(probe.starts.load(Ordering::SeqCst), 0); + assert!( + warning_rx + .recv() + .await + .is_some_and(|warning| warning.contains("whole-response turn-end gate")) + ); + + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + let (output, _history) = run.await.unwrap(); + assert!(matches!(output.result, EngineRunExit::Paused)); + assert_eq!(probe.starts.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_pause() { + struct BlockingEarlyPolicy { + entered: Arc, + release: Arc, + } + + #[async_trait] + impl Interceptor for BlockingEarlyPolicy { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn pre_tool_call( + &self, + _info: &mut ToolCallInfo<'_>, + ) -> InterceptorResult { + self.entered.notify_one(); + self.release.notified().await; + Ok(PreToolAction::Continue) + } + } + + let (client, tx) = BoundedStreamClient::new(1); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let tool = SlowTool::new("blocked_policy_tool", 0); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(BlockingEarlyPolicy { + entered: entered.clone(), + release, + }); + engine.register_tool(tool.definition()); + let pause = engine.pause_sender(); + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + + tx.send(Ok(Event::tool_use_start( + 0, + "call_blocked_policy", + "blocked_policy_tool", + ))) + .await + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))) + .await + .unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), entered.notified()) + .await + .expect("pre-tool policy entered"); + + tx.send(Ok(Event::ping())).await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), tx.send(Ok(Event::ping()))) + .await + .expect("background pump continues draining the bounded provider stream") + .unwrap(); + pause.send(()).await.unwrap(); + drop(tx); + + let (output, _history) = run.await.unwrap(); + assert!(matches!(output.result, EngineRunExit::Paused)); + assert_eq!(probe.call_count(), 0); +} + #[tokio::test] async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_result() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); @@ -457,6 +633,146 @@ async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_ assert_eq!((calls, results), (2, 2)); } +#[tokio::test] +async fn early_terminal_result_commits_while_response_stream_is_open() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = BarrierTool::new("commit_while_streaming"); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_max_turns(Some(1)); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(tool.definition()); + let (committed_tx, mut committed_rx) = tokio::sync::mpsc::unbounded_channel(); + engine.on_history_append(move |item| { + if matches!(item, Item::ToolResult { call_id, .. } if call_id == "call_commit_open") { + let _ = committed_tx.send(()); + } + Ok(()) + }); + + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + tx.send(Ok(Event::tool_use_start( + 0, + "call_commit_open", + "commit_while_streaming", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tokio::time::timeout(Duration::from_secs(1), probe.started.notified()) + .await + .expect("tool starts while response remains open"); + probe.release.notify_one(); + tokio::time::timeout(Duration::from_secs(1), committed_rx.recv()) + .await + .expect("terminal result is durably committed before response release") + .expect("commit notification channel remains open"); + + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + let _ = run.await.unwrap(); +} + +#[tokio::test] +async fn reverse_stop_order_preserves_model_call_indexes() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let contexts = Arc::new(Mutex::new(Vec::new())); + let mut engine = Engine::new(client); + engine.set_max_turns(Some(1)); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(ContextRecordingTool::new("ordered_a", contexts.clone()).definition()); + engine.register_tool(ContextRecordingTool::new("ordered_b", contexts.clone()).definition()); + + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + tx.send(Ok(Event::tool_use_start(0, "call_order_a", "ordered_a"))) + .unwrap(); + tx.send(Ok(Event::tool_use_start(1, "call_order_b", "ordered_b"))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(1, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(1))).unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + let _ = run.await.unwrap(); + + let by_id = contexts + .lock() + .unwrap() + .iter() + .map(|context| (context.call_id.clone(), context.call_index)) + .collect::>(); + assert_eq!(by_id.get("call_order_a"), Some(&0)); + assert_eq!(by_id.get("call_order_b"), Some(&1)); +} + +#[tokio::test] +async fn reverse_stop_order_projects_next_provider_request_in_model_order() { + let first_response = vec![ + Event::tool_use_start(0, "call_project_a", "project_a"), + Event::tool_use_start(1, "call_project_b", "project_b"), + Event::tool_input_delta(1, r#"{}"#), + Event::tool_use_stop(1), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]; + let second_response = vec![ + Event::text_delta(0, "done"), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]; + let client = MockLlmClient::with_responses(vec![first_response, second_response]) + .with_completion_support(ToolCallCompletionSupport::PerBlock); + let probe = client.clone(); + let contexts = Arc::new(Mutex::new(Vec::new())); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(ContextRecordingTool::new("project_a", contexts.clone()).definition()); + engine.register_tool(ContextRecordingTool::new("project_b", contexts).definition()); + let mut history = History::new(); + + let _ = engine.run(&mut history, "run").await; + + let requests = probe.requests(); + assert_eq!(requests.len(), 2); + let tool_order = requests[1] + .items + .iter() + .filter_map(|item| match item { + Item::ToolCall { call_id, .. } => Some(format!("call:{call_id}")), + Item::ToolResult { call_id, .. } => Some(format!("result:{call_id}")), + _ => None, + }) + .collect::>(); + assert_eq!( + tool_order, + [ + "call:call_project_a", + "call:call_project_b", + "result:call_project_a", + "result:call_project_b", + ] + ); +} + #[tokio::test] async fn later_history_failure_terminalizes_already_started_early_sibling() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); @@ -525,6 +841,108 @@ async fn later_history_failure_terminalizes_already_started_early_sibling() { ))); } +#[tokio::test] +async fn persisted_started_call_is_closed_unknown_and_never_resumed() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = SlowTool::new("must_not_resume_started", 0); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_max_turns(Some(1)); + engine.register_tool(tool.definition()); + let mut history = History::from_items(vec![ + Item::user_message("run"), + Item::tool_call_json( + "call_started_before_restore", + "must_not_resume_started", + serde_json::json!({}), + ) + .with_tool_execution_metadata( + 0, + Some("restored-batch:call_started_before_restore".to_string()), + ) + .with_status(agen::llm_client::ItemStatus::InProgress), + ]); + + let run = tokio::spawn(async move { + let output = engine.resume(&mut history).await; + (output, history) + }); + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + let (_output, history) = run.await.unwrap(); + + assert_eq!(probe.call_count(), 0); + assert!(history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call_started_before_restore" + ))); +} + +#[tokio::test] +async fn abort_during_early_admission_commits_synthetic_terminal() { + struct AbortAdmission; + + #[async_trait] + impl Interceptor for AbortAdmission { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn pre_tool_call( + &self, + _info: &mut ToolCallInfo<'_>, + ) -> InterceptorResult { + Ok(PreToolAction::Abort("denied after completion".to_string())) + } + } + + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let tool = SlowTool::new("must_not_run_aborted", 0); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(AbortAdmission); + engine.register_tool(tool.definition()); + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + + tx.send(Ok(Event::tool_use_start( + 0, + "call_aborted_admission", + "must_not_run_aborted", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + drop(tx); + let (output, history) = run.await.unwrap(); + + assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(EngineError::Aborted(ref reason))) + if reason == "denied after completion" + )); + assert_eq!(probe.call_count(), 0); + assert!(history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { + call_id, + summary, + .. + } if call_id == "call_aborted_admission" && summary.contains("denied after completion") + ))); +} + #[tokio::test] async fn malformed_completed_arguments_commit_error_without_executing_tool() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); diff --git a/crates/session-store/src/logged_item.rs b/crates/session-store/src/logged_item.rs index f6547265b..6855b0973 100644 --- a/crates/session-store/src/logged_item.rs +++ b/crates/session-store/src/logged_item.rs @@ -7,13 +7,14 @@ //! save / replay boundaries. //! //! Fields kept here are limited to what is needed to reconstruct a worker -//! `Item` for replay. `id` and `status` annotations are intentionally dropped -//! (they are output-side metadata; replayed items synthesize fresh `None`). +//! `Item` for replay. `id` annotations are intentionally dropped because they +//! are output-side metadata. ToolCall execution status/order are retained because +//! they fence durable resume from replaying an already-started side effect. //! `Reasoning::encrypted_content` is preserved because OpenAI Responses ZDR //! requires it on stateless re-send. use agen::{ - llm_client::types::{ContentPart, Item, Role}, + llm_client::types::{ContentPart, Item, ItemStatus, Role}, tool::{Attachment, ImageAttachment, ToolResultDisposition}, }; use base64::{Engine as _, engine::general_purpose::STANDARD}; @@ -53,6 +54,12 @@ pub enum LoggedItem { call_id: String, name: String, arguments: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + call_index: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + execution_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + status: Option, }, ToolResult { call_id: String, @@ -119,11 +126,17 @@ impl From<&Item> for LoggedItem { call_id, name, arguments, + call_index, + execution_id, + status, .. } => Self::ToolCall { call_id: call_id.clone(), name: name.clone(), arguments: arguments.clone(), + call_index: *call_index, + execution_id: execution_id.clone(), + status: *status, }, Item::ToolResult { call_id, @@ -176,12 +189,17 @@ impl From for Item { call_id, name, arguments, + call_index, + execution_id, + status, } => Item::ToolCall { id: None, call_id, name, arguments, - status: None, + call_index, + execution_id, + status, }, LoggedItem::ToolResult { call_id, @@ -327,7 +345,9 @@ mod tests { #[test] fn round_trip_tool_call() { - let original = Item::tool_call("call_42", "get_weather", r#"{"city":"Tokyo"}"#); + let original = Item::tool_call("call_42", "get_weather", r#"{"city":"Tokyo"}"#) + .with_tool_execution_metadata(3, Some("batch:call_42".to_string())) + .with_status(ItemStatus::InProgress); let logged: LoggedItem = (&original).into(); let json = serde_json::to_string(&logged).unwrap(); let parsed: LoggedItem = serde_json::from_str(&json).unwrap(); @@ -336,11 +356,17 @@ mod tests { call_id, name, arguments, + call_index, + execution_id, + status, .. } => { assert_eq!(call_id, "call_42"); assert_eq!(name, "get_weather"); assert_eq!(arguments, r#"{"city":"Tokyo"}"#); + assert_eq!(call_index, Some(3)); + assert_eq!(execution_id.as_deref(), Some("batch:call_42")); + assert_eq!(status, Some(ItemStatus::InProgress)); } other => panic!("unexpected variant: {other:?}"), } diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index 752187064..72ec0a140 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -356,6 +356,7 @@ fn project_item(item: &LoggedItem) -> Option { call_id, name, arguments, + .. } => Some(SessionSnapshotEntryData::ToolCall { call_id: call_id.clone(), name: name.clone(), diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 802d0c7d7..2b403cd39 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -363,6 +363,10 @@ impl WorkerInterceptor { #[async_trait] impl Interceptor for WorkerInterceptor { + fn supports_early_tool_dispatch(&self) -> bool { + self.registry.on_turn_end.is_empty() + } + async fn on_prompt_submit( &self, context: PromptSubmitContext<'_, SessionHistoryMetadata>, diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 8cac5d7e5..f0b05eeec 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -1371,6 +1371,9 @@ mod tests { call_id: format!("call-{index}"), name: name.to_string(), arguments: "{}".to_string(), + call_index: None, + execution_id: None, + status: None, }, ), }); diff --git a/docs/manifest.toml b/docs/manifest.toml index 274a792e8..c1bcccc9f 100644 --- a/docs/manifest.toml +++ b/docs/manifest.toml @@ -123,8 +123,9 @@ ref = "anthropic/claude-sonnet-4-6" # 任意。デフォルト: "after_response"。 # "on_tool_call_complete" は provider が個別 ToolCall の完成を保証する場合だけ、 -# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider では -# 応答完了まで明示的に遅延する。これは直列実行モードではなく、先行 call の完了を +# そのブロック完成後に実行開始を前倒しする。安全な完成境界がない provider、または +# 応答全体を見て実行前 Pause を判断する turn-end hook がある場合は、応答完了まで +# 警告付きで明示的に遅延する。これは直列実行モードではなく、先行 call の完了を # 待たず後続 call も開始するため、完了順・DB 反映順・採番順は保証しない。 # 早期開始後に stream が切断された場合は自動 continuation せず、安全に terminalize して停止する。 # tool_call_dispatch = "on_tool_call_complete" From dd2bc41ab87a3981132a136385bb5d2c7a4d4b77 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 12:03:02 +0900 Subject: [PATCH 04/11] fix: harden external command lifecycle --- crates/workdir/src/external.rs | 56 +- crates/workspace-server/src/server.rs | 707 +++++++++++++++++++++++--- crates/yoi/src/workdir_share.rs | 25 +- 3 files changed, 700 insertions(+), 88 deletions(-) diff --git a/crates/workdir/src/external.rs b/crates/workdir/src/external.rs index a1b12ec06..5557231ff 100644 --- a/crates/workdir/src/external.rs +++ b/crates/workdir/src/external.rs @@ -352,6 +352,27 @@ impl<'de> Deserialize<'de> for ExternalWorkdirOperationResult { } } +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ExternalWorkdirOperationError { + code: crate::http::WorkdirTransportErrorCode, +} + +impl ExternalWorkdirOperationError { + /// Collapse provider failures to a closed error code. Provider-authored + /// diagnostics never cross the External Workdir trust boundary. + pub fn from_transport_error(error: WorkdirTransportError) -> Self { + Self { code: error.code } + } + + pub fn into_transport_error(self) -> WorkdirTransportError { + WorkdirTransportError { + code: self.code, + message: format!("External Workdir provider reported {}", self.code.as_str()), + } + } +} + #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(tag = "outcome", rename_all = "snake_case", deny_unknown_fields)] pub enum ExternalWorkdirOperationOutcome { @@ -359,7 +380,7 @@ pub enum ExternalWorkdirOperationOutcome { result: ExternalWorkdirOperationResult, }, Failed { - error: WorkdirTransportError, + error: ExternalWorkdirOperationError, }, Cancelled, } @@ -557,6 +578,39 @@ mod tests { } } + #[test] + fn provider_failures_are_closed_and_drop_provider_authored_diagnostics() { + let error = ExternalWorkdirOperationError::from_transport_error(WorkdirTransportError { + code: crate::http::WorkdirTransportErrorCode::Internal, + message: "/home/operator/private provider-secret\ncontrol".to_string(), + }); + let frame = ExternalWorkdirProviderFrame::current( + ExternalWorkdirProviderMessage::OperationResult { + generation: 1, + operation_id: ExternalWorkdirOperationId::new("operation-failed").unwrap(), + outcome: ExternalWorkdirOperationOutcome::Failed { + error: error.clone(), + }, + }, + ); + let value = serde_json::to_value(&frame).unwrap(); + assert_no_provider_authority(&value); + assert!(value.to_string().len() < 512); + assert_eq!( + error.into_transport_error().message, + "External Workdir provider reported internal" + ); + + let provider_message = serde_json::to_string(&frame).unwrap().replace( + r#""code":"internal""#, + r#""code":"internal","message":"provider-secret""#, + ); + assert!( + serde_json::from_str::(&provider_message).is_err(), + "provider-authored diagnostics must not be accepted" + ); + } + #[test] fn unsupported_protocol_versions_and_invalid_read_limits_fail_closed() { let future = r#"{ diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 57f68cb53..f70330a97 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -758,6 +758,7 @@ struct ExternalProviderConnection { expires_at: Option>, capabilities: workdir::WorkdirSessionCapabilities, admission: Arc, + shutdown_confirmed: Arc, sender: tokio::sync::mpsc::Sender, } @@ -985,6 +986,68 @@ impl ExternalProviderWorkdirSession { } } + async fn provider_command_handle( + &self, + external_handle: &workdir::CommandHandle, + ) -> std::result::Result { + if self.closed.load(Ordering::Acquire) { + return Err(workdir::WorkdirError::SessionClosed); + } + self.active_commands + .lock() + .await + .get(&external_handle.0) + .cloned() + .ok_or_else(|| workdir::WorkdirError::UnknownCommand(external_handle.0.clone())) + } + + async fn remove_active_command_if_matches( + &self, + external_handle: &workdir::CommandHandle, + provider_handle: &workdir::CommandHandle, + ) { + let mut commands = self.active_commands.lock().await; + if commands.get(&external_handle.0) == Some(provider_handle) { + commands.remove(&external_handle.0); + } + } + + async fn terminate_provider_command( + &self, + provider_handle: workdir::CommandHandle, + ) -> 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. + let _ = self + .operate(WorkdirSessionOperation::CommandCancel( + provider_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")), + } + } + fn mismatch(expected: &str) -> workdir::WorkdirError { workdir::WorkdirError::Unavailable(format!( "External Workdir provider returned a mismatched result; expected {expected}" @@ -1146,23 +1209,58 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { .operate(WorkdirSessionOperation::CommandStart(request)) .await? { - WorkdirSessionOperationResult::CommandStart(handle) => { + WorkdirSessionOperationResult::CommandStart(provider_handle) => { + let duplicate_handles = { + let commands = self.active_commands.lock().await; + commands + .iter() + .filter(|(_, active)| *active == &provider_handle) + .map(|(external, _)| external.clone()) + .collect::>() + }; + if !duplicate_handles.is_empty() { + // A provider handle must identify exactly one command in one + // generation. Remove every ambiguous route before attempting + // cleanup so no caller can address the wrong process. + let mut commands = self.active_commands.lock().await; + for external_handle in duplicate_handles { + commands.remove(&external_handle); + } + let quarantine_handle = + workdir::CommandHandle(format!("quarantined-command-{}", Uuid::now_v7())); + commands.insert(quarantine_handle.0.clone(), provider_handle.clone()); + drop(commands); + if self + .terminate_provider_command(provider_handle.clone()) + .await + .is_ok() + { + self.remove_active_command_if_matches(&quarantine_handle, &provider_handle) + .await; + } + return Err(workdir::WorkdirError::Unavailable( + "External Workdir provider reused an active command handle".to_string(), + )); + } + let external_handle = + workdir::CommandHandle(format!("external-command-{}", Uuid::now_v7())); self.active_commands .lock() .await - .insert(handle.0.clone(), handle.clone()); - Ok(handle) + .insert(external_handle.0.clone(), provider_handle); + Ok(external_handle) } _ => Err(Self::mismatch("command_start")), } } async fn command_status( &self, - handle: workdir::CommandHandle, + external_handle: workdir::CommandHandle, ) -> std::result::Result { self.ensure_capability(workdir::WorkdirSessionCapability::Command)?; + let provider_handle = self.provider_command_handle(&external_handle).await?; match self - .operate(WorkdirSessionOperation::CommandStatus(handle)) + .operate(WorkdirSessionOperation::CommandStatus(provider_handle)) .await? { WorkdirSessionOperationResult::CommandStatus(status) => Ok(status), @@ -1171,10 +1269,12 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { } async fn command_output( &self, - request: workdir::CommandOutputRequest, + mut request: workdir::CommandOutputRequest, ) -> std::result::Result { self.ensure_capability(workdir::WorkdirSessionCapability::Command)?; - let command_id = request.handle.0.clone(); + let external_handle = request.handle.clone(); + let provider_handle = self.provider_command_handle(&external_handle).await?; + request.handle = provider_handle.clone(); match self .operate(WorkdirSessionOperation::CommandOutput(request)) .await? @@ -1182,7 +1282,8 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { WorkdirSessionOperationResult::CommandOutput(output) => { if output.status != workdir::CommandStatus::Running && output.next_cursor.is_none() { - self.active_commands.lock().await.remove(&command_id); + self.remove_active_command_if_matches(&external_handle, &provider_handle) + .await; } Ok(output) } @@ -1191,11 +1292,12 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { } async fn cancel_command( &self, - handle: workdir::CommandHandle, + external_handle: workdir::CommandHandle, ) -> std::result::Result<(), workdir::WorkdirError> { self.ensure_capability(workdir::WorkdirSessionCapability::Command)?; + let provider_handle = self.provider_command_handle(&external_handle).await?; match self - .operate(WorkdirSessionOperation::CommandCancel(handle)) + .operate(WorkdirSessionOperation::CommandCancel(provider_handle)) .await? { WorkdirSessionOperationResult::CommandCancel => Ok(()), @@ -1207,31 +1309,50 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { if self.closed.load(Ordering::Acquire) { return Ok(()); } + if self.connection.shutdown_confirmed.load(Ordering::Acquire) { + self.active_commands.lock().await.clear(); + self.closed.store(true, Ordering::Release); + return Ok(()); + } let commands = self .active_commands .lock() .await - .drain() - .map(|(_, handle)| handle) + .iter() + .map(|(external, provider)| { + (workdir::CommandHandle(external.clone()), provider.clone()) + }) .collect::>(); - for handle in commands { + let mut cleanup_errors = Vec::new(); + for (external_handle, provider_handle) in commands { // Detach/session replacement is terminal authority for commands - // started through this exact provider generation. Request cancel, - // then wait for provider-owned process reaping before closing. - let _ = self - .operate(WorkdirSessionOperation::CommandCancel(handle.clone())) - .await; - let _ = self - .operate(WorkdirSessionOperation::CommandOutput( - workdir::CommandOutputRequest { - handle, - cursor: 0, - limit: workdir::external::MAX_EXTERNAL_COMMAND_OUTPUT_BYTES, - wait: true, - }, - )) - .await; + // 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()) + .await + { + Ok(()) => { + self.remove_active_command_if_matches(&external_handle, &provider_handle) + .await; + } + Err(error) => { + let still_active = self.active_commands.lock().await.get(&external_handle.0) + == Some(&provider_handle); + if still_active { + cleanup_errors.push(error.to_string()); + } + } + } } + if !cleanup_errors.is_empty() && !self.connection.shutdown_confirmed.load(Ordering::Acquire) + { + return Err(workdir::WorkdirError::Unavailable(format!( + "External Workdir command cleanup was not confirmed: {}", + cleanup_errors.join("; ") + ))); + } + self.active_commands.lock().await.clear(); self.closed.store(true, Ordering::Release); Ok(()) } @@ -17097,6 +17218,25 @@ fn cancel_external_workdir_expiry(api: &WorkspaceApi, grant_id: &str) { } } +async fn stop_external_workdir_provider( + api: &WorkspaceApi, + grant: &ExternalWorkdirGrantRecord, + reason: workdir::external::ExternalWorkdirRevokeReason, +) { + let connection = { + api.external_workdir_providers + .lock() + .expect("External Workdir provider registry poisoned") + .disconnect(&grant.grant_id, grant.generation) + }; + if let Some(connection) = connection { + let _ = connection + .sender + .send(ExternalProviderCommand::Revoke { reason }) + .await; + } +} + async fn fence_external_workdir_grant_expired( api: &WorkspaceApi, grant_id: &str, @@ -17127,20 +17267,6 @@ async fn fence_external_workdir_grant_expired( // A reconnect may advance generation between the read and fence. continue; } - let connection = { - api.external_workdir_providers - .lock() - .expect("External Workdir provider registry poisoned") - .disconnect(&grant.grant_id, grant.generation) - }; - if let Some(connection) = connection { - let _ = connection - .sender - .send(ExternalProviderCommand::Revoke { - reason: workdir::external::ExternalWorkdirRevokeReason::Expired, - }) - .await; - } api.store.record_external_workdir_audit( &grant.workspace_id, "system", @@ -17166,7 +17292,14 @@ async fn expire_external_workdir_grant(api: &WorkspaceApi, grant_id: &str) -> Ap if !matches!(grant.status.as_str(), "revoked" | "expired") { return Ok(false); } - cleanup_external_workdir_attachments(api, &grant.workdir_id).await?; + let cleanup = cleanup_external_workdir_attachments(api, &grant.workdir_id).await; + stop_external_workdir_provider( + api, + &grant, + workdir::external::ExternalWorkdirRevokeReason::Expired, + ) + .await; + cleanup?; Ok(true) } @@ -17412,6 +17545,7 @@ async fn serve_external_workdir_provider( expires_at, capabilities: registration.capabilities, admission: Arc::new(tokio::sync::Semaphore::new(16)), + shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)), sender, }); if api @@ -17464,6 +17598,7 @@ async fn serve_external_workdir_provider( heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); let mut heartbeat_sequence = 0_u64; let mut last_heartbeat_ack = tokio::time::Instant::now(); + let mut revoking = false; let terminal_status = loop { tokio::select! { _ = heartbeat.tick() => { @@ -17481,6 +17616,13 @@ async fn serve_external_workdir_provider( } command = commands.recv() => match command { Some(ExternalProviderCommand::Operation { operation_id, operation, response }) => { + if revoking { + let _ = response.send(Err(WorkdirTransportError { + code: workdir::http::WorkdirTransportErrorCode::Unavailable, + message: "External Workdir provider shutdown is in progress".to_string(), + })); + continue; + } if pending.len() >= 16 { let _ = response.send(Err(WorkdirTransportError { code: workdir::http::WorkdirTransportErrorCode::Unavailable, @@ -17521,12 +17663,17 @@ async fn serve_external_workdir_provider( } } Some(ExternalProviderCommand::Revoke { reason }) => { + if revoking { + continue; + } let frame = ExternalWorkdirServerFrame::current(ExternalWorkdirServerMessage::Revoke { generation: connection.generation, reason, }); - let _ = send_external_workdir_server_frame(&mut socket, &frame).await; - break "revoked"; + if !send_external_workdir_server_frame(&mut socket, &frame).await { + break "offline"; + } + revoking = true; } None => break "offline", }, @@ -17540,7 +17687,9 @@ async fn serve_external_workdir_provider( if let Some(response) = pending.remove(operation_id.as_str()) { let result = match outcome { ExternalWorkdirOperationOutcome::Completed { result } => Ok(result.into_inner()), - ExternalWorkdirOperationOutcome::Failed { error } => Err(error), + ExternalWorkdirOperationOutcome::Failed { error } => { + Err(error.into_transport_error()) + } ExternalWorkdirOperationOutcome::Cancelled => Err(WorkdirTransportError { code: workdir::http::WorkdirTransportErrorCode::Unavailable, message: "External Workdir operation was cancelled".to_string(), @@ -17554,7 +17703,10 @@ async fn serve_external_workdir_provider( last_heartbeat_ack = tokio::time::Instant::now(); } ExternalWorkdirProviderMessage::RevokeAcknowledged { generation } - if generation == connection.generation => break "revoked", + if revoking && generation == connection.generation => { + connection.shutdown_confirmed.store(true, Ordering::Release); + break "revoked"; + } _ => break "offline", } } @@ -17658,7 +17810,14 @@ async fn scoped_get_external_workdir_grant( // Terminal cleanup is retryable: a prior expiry/revoke may have committed // its fence before Runtime/session projection completed. if matches!(grant.status.as_str(), "expired" | "revoked") { - cleanup_external_workdir_attachments(&api, &grant.workdir_id).await?; + let cleanup = cleanup_external_workdir_attachments(&api, &grant.workdir_id).await; + let reason = if grant.status == "expired" { + workdir::external::ExternalWorkdirRevokeReason::Expired + } else { + workdir::external::ExternalWorkdirRevokeReason::UserRequested + }; + stop_external_workdir_provider(&api, &grant, reason).await; + cleanup?; cancel_external_workdir_expiry(&api, &grant_id); } Ok(Json(external_workdir_grant_response(&grant)?)) @@ -17690,20 +17849,6 @@ async fn scoped_revoke_external_workdir_grant( "revoked", &updated_at, )?; - let connection = { - api.external_workdir_providers - .lock() - .expect("External Workdir provider registry poisoned") - .disconnect(&grant_id, grant.generation) - }; - if let Some(connection) = connection { - let _ = connection - .sender - .send(ExternalProviderCommand::Revoke { - reason: workdir::external::ExternalWorkdirRevokeReason::UserRequested, - }) - .await; - } grant.status = "revoked".to_string(); grant.updated_at = updated_at; api.store.record_external_workdir_audit( @@ -17723,7 +17868,14 @@ async fn scoped_revoke_external_workdir_grant( } // Cleanup is deliberately retried for terminal grants. A prior revoke may // have committed its fence before a Runtime/session projection failure. - cleanup_external_workdir_attachments(&api, &grant.workdir_id).await?; + let cleanup = cleanup_external_workdir_attachments(&api, &grant.workdir_id).await; + stop_external_workdir_provider( + &api, + &grant, + workdir::external::ExternalWorkdirRevokeReason::UserRequested, + ) + .await; + cleanup?; cancel_external_workdir_expiry(&api, &grant_id); Ok(Json(external_workdir_grant_response(&grant)?)) } @@ -26338,6 +26490,7 @@ mod tests { expires_at: Some(Utc::now() + Duration::minutes(1)), capabilities: workdir::WorkdirSessionCapabilities::READ_ONLY, admission: Arc::new(tokio::sync::Semaphore::new(16)), + shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)), sender, }); let session = Arc::new(ExternalProviderWorkdirSession::new(connection, None)); @@ -26402,6 +26555,7 @@ mod tests { expires_at: None, capabilities: workdir::WorkdirSessionCapabilities::READ_WRITE, admission: Arc::new(tokio::sync::Semaphore::new(16)), + shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)), sender, }); let session = Arc::new(ExternalProviderWorkdirSession::new(connection, None)); @@ -26494,6 +26648,7 @@ mod tests { expires_at: None, capabilities, admission: Arc::new(tokio::sync::Semaphore::new(16)), + shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)), sender, }); let session = Arc::new(ExternalProviderWorkdirSession::new(connection, None)); @@ -26543,13 +26698,19 @@ mod tests { .. }) )); - let handle = workdir::CommandHandle("command-1".to_string()); + let provider_handle = workdir::CommandHandle("command-1".to_string()); response .send(Ok(WorkdirSessionOperationResult::CommandStart( - handle.clone(), + provider_handle.clone(), ))) .unwrap(); - assert_eq!(starter.await.unwrap().unwrap(), handle); + let external_handle = starter.await.unwrap().unwrap(); + assert_ne!(external_handle, provider_handle); + assert!(matches!( + session.command_status(provider_handle.clone()).await, + Err(workdir::WorkdirError::UnknownCommand(_)) + )); + assert!(commands.try_recv().is_err()); let closer = { let session = session.clone(); @@ -26565,7 +26726,7 @@ mod tests { }; assert!(matches!( operation, - WorkdirSessionOperation::CommandCancel(ref current) if current == &handle + WorkdirSessionOperation::CommandCancel(ref current) if current == &provider_handle )); response .send(Ok(WorkdirSessionOperationResult::CommandCancel)) @@ -26585,6 +26746,85 @@ 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, + }, + ))) + .unwrap(); + assert!(matches!( + closer.await.unwrap(), + Err(workdir::WorkdirError::Unavailable(_)) + )); + + let status = { + let session = session.clone(); + let handle = external_handle.clone(); + tokio::spawn(async move { session.command_status(handle).await }) + }; + let Some(ExternalProviderCommand::Operation { + operation, + response, + .. + }) = commands.recv().await + else { + panic!("expected retained command status after failed close"); + }; + assert!(matches!( + operation, + WorkdirSessionOperation::CommandStatus(ref current) if current == &provider_handle + )); + response + .send(Ok(WorkdirSessionOperationResult::CommandStatus( + workdir::CommandStatus::Running, + ))) + .unwrap(); + assert_eq!( + status.await.unwrap().unwrap(), + workdir::CommandStatus::Running + ); + + let closer = { + let session = session.clone(); + tokio::spawn(async move { session.close().await }) + }; + let Some(ExternalProviderCommand::Operation { + operation, + response, + .. + }) = commands.recv().await + else { + panic!("expected retried command cancellation"); + }; + assert!(matches!( + operation, + WorkdirSessionOperation::CommandCancel(ref current) if current == &provider_handle + )); + response + .send(Ok(WorkdirSessionOperationResult::CommandCancel)) + .unwrap(); + let Some(ExternalProviderCommand::Operation { + operation, + response, + .. + }) = commands.recv().await + else { + panic!("expected retried terminal command output"); + }; + assert!(matches!( + operation, + WorkdirSessionOperation::CommandOutput(workdir::CommandOutputRequest { + wait: true, + .. + }) + )); response .send(Ok(WorkdirSessionOperationResult::CommandOutput( workdir::CommandOutput { @@ -26600,11 +26840,211 @@ mod tests { .unwrap(); closer.await.unwrap().unwrap(); assert!(matches!( - session.command_status(handle).await, + session.command_status(external_handle).await, Err(workdir::WorkdirError::SessionClosed) )); } + #[tokio::test] + async fn external_provider_reused_command_handle_invalidates_ambiguous_routes() { + let (sender, mut commands) = tokio::sync::mpsc::channel(8); + let connection = Arc::new(ExternalProviderConnection { + grant_id: "grant-duplicate-command".to_string(), + workdir_id: "external-duplicate-command".to_string(), + provider_instance_id: "provider-duplicate-command".to_string(), + generation: 3, + expires_at: None, + capabilities: workdir::WorkdirSessionCapabilities::COMMAND_ONLY, + admission: Arc::new(tokio::sync::Semaphore::new(16)), + shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)), + sender, + }); + let session = Arc::new(ExternalProviderWorkdirSession::new(connection, None)); + let request = || workdir::CommandRequest { + command: "sleep 30".to_string(), + timeout_secs: 30, + output_limit: 1024, + cwd: workdir::WorkdirPath::root(), + spill_dir: None, + tool_call_id: None, + }; + let provider_handle = workdir::CommandHandle("command-reused".to_string()); + + let first = { + let session = session.clone(); + let request = request(); + tokio::spawn(async move { session.start_command(request).await }) + }; + let Some(ExternalProviderCommand::Operation { response, .. }) = commands.recv().await + else { + panic!("expected first command start"); + }; + response + .send(Ok(WorkdirSessionOperationResult::CommandStart( + provider_handle.clone(), + ))) + .unwrap(); + let first_external = first.await.unwrap().unwrap(); + + let duplicate = { + let session = session.clone(); + let request = request(); + tokio::spawn(async move { session.start_command(request).await }) + }; + let Some(ExternalProviderCommand::Operation { response, .. }) = commands.recv().await + else { + panic!("expected duplicate command start"); + }; + response + .send(Ok(WorkdirSessionOperationResult::CommandStart( + provider_handle.clone(), + ))) + .unwrap(); + let Some(ExternalProviderCommand::Operation { + operation, + response, + .. + }) = commands.recv().await + else { + panic!("expected ambiguous command cancellation"); + }; + assert!(matches!( + operation, + WorkdirSessionOperation::CommandCancel(ref current) if current == &provider_handle + )); + response + .send(Ok(WorkdirSessionOperationResult::CommandCancel)) + .unwrap(); + let Some(ExternalProviderCommand::Operation { response, .. }) = commands.recv().await + else { + panic!("expected ambiguous command terminal output"); + }; + 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, + }, + ))) + .unwrap(); + assert!(matches!( + duplicate.await.unwrap(), + Err(workdir::WorkdirError::Unavailable(_)) + )); + assert!(matches!( + session.command_status(first_external).await, + Err(workdir::WorkdirError::UnknownCommand(_)) + )); + assert!(commands.try_recv().is_err()); + session.close().await.unwrap(); + } + + #[tokio::test] + async fn external_provider_generation_sessions_reject_each_others_command_handles() { + let (first_sender, mut first_commands) = tokio::sync::mpsc::channel(4); + let first = Arc::new(ExternalProviderWorkdirSession::new( + Arc::new(ExternalProviderConnection { + grant_id: "grant-generation-command".to_string(), + workdir_id: "external-generation-command".to_string(), + provider_instance_id: "provider-generation-command".to_string(), + generation: 1, + expires_at: None, + capabilities: workdir::WorkdirSessionCapabilities::COMMAND_ONLY, + admission: Arc::new(tokio::sync::Semaphore::new(16)), + shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)), + sender: first_sender, + }), + None, + )); + let (second_sender, mut second_commands) = tokio::sync::mpsc::channel(4); + let second = Arc::new(ExternalProviderWorkdirSession::new( + Arc::new(ExternalProviderConnection { + grant_id: "grant-generation-command".to_string(), + workdir_id: "external-generation-command".to_string(), + provider_instance_id: "provider-generation-command".to_string(), + generation: 2, + expires_at: None, + capabilities: workdir::WorkdirSessionCapabilities::COMMAND_ONLY, + admission: Arc::new(tokio::sync::Semaphore::new(16)), + shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)), + sender: second_sender, + }), + None, + )); + let request = || workdir::CommandRequest { + command: "sleep 30".to_string(), + timeout_secs: 30, + output_limit: 1024, + cwd: workdir::WorkdirPath::root(), + spill_dir: None, + tool_call_id: None, + }; + let provider_handle = workdir::CommandHandle("command-1".to_string()); + + let first_start = { + let first = first.clone(); + let request = request(); + tokio::spawn(async move { first.start_command(request).await }) + }; + let Some(ExternalProviderCommand::Operation { response, .. }) = first_commands.recv().await + else { + panic!("expected generation-one command start"); + }; + response + .send(Ok(WorkdirSessionOperationResult::CommandStart( + provider_handle.clone(), + ))) + .unwrap(); + let stale_handle = first_start.await.unwrap().unwrap(); + + let second_start = { + let second = second.clone(); + let request = request(); + tokio::spawn(async move { second.start_command(request).await }) + }; + let Some(ExternalProviderCommand::Operation { response, .. }) = + second_commands.recv().await + else { + panic!("expected generation-two command start"); + }; + response + .send(Ok(WorkdirSessionOperationResult::CommandStart( + provider_handle, + ))) + .unwrap(); + let current_handle = second_start.await.unwrap().unwrap(); + assert_ne!(stale_handle, current_handle); + assert!(matches!( + second.command_status(stale_handle).await, + Err(workdir::WorkdirError::UnknownCommand(_)) + )); + assert!(second_commands.try_recv().is_err()); + + let status = { + let second = second.clone(); + tokio::spawn(async move { second.command_status(current_handle).await }) + }; + let Some(ExternalProviderCommand::Operation { response, .. }) = + second_commands.recv().await + else { + panic!("expected current-generation command status"); + }; + response + .send(Ok(WorkdirSessionOperationResult::CommandStatus( + workdir::CommandStatus::Running, + ))) + .unwrap(); + assert_eq!( + status.await.unwrap().unwrap(), + workdir::CommandStatus::Running + ); + } + #[tokio::test] async fn backend_expiry_fences_offline_grant_and_persists_cleanup_completion() { let dir = tempfile::tempdir().unwrap(); @@ -26979,6 +27419,7 @@ mod tests { expires_at: Some(expires_at), capabilities: workdir::WorkdirSessionCapabilities::READ_ONLY, admission: Arc::new(tokio::sync::Semaphore::new(16)), + shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)), sender, }); assert!( @@ -27406,7 +27847,8 @@ mod tests { tokio::spawn(async move { broker_session .start_command(workdir::CommandRequest { - command: "printf command-ok".to_string(), + command: "sleep 30 & child=$!; echo $child > child.pid; wait $child" + .to_string(), timeout_secs: 30, output_limit: 1024, cwd: workdir::WorkdirPath::root(), @@ -27459,6 +27901,14 @@ mod tests { .unwrap(); let command_handle = command_start.await.unwrap().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 = { let broker_session = broker_session.clone(); let handle = command_handle.clone(); @@ -27468,7 +27918,7 @@ mod tests { handle, cursor: 0, limit: 1024, - wait: true, + wait: false, }) .await }) @@ -27487,7 +27937,10 @@ mod tests { }; assert!(matches!( operation.as_inner(), - WorkdirSessionOperation::CommandOutput(_) + WorkdirSessionOperation::CommandOutput(workdir::CommandOutputRequest { + wait: false, + .. + }) )); let result = workdir::http::dispatch_workdir_session_operation( &local_session, @@ -27512,17 +27965,106 @@ mod tests { .await .unwrap(); let command_output = command_output.await.unwrap().unwrap(); - assert_eq!(command_output.content, "command-ok"); - assert_eq!(command_output.exit_code, Some(0)); + assert_eq!(command_output.status, workdir::CommandStatus::Running); assert_eq!(command_output.output_path, None); - let Json(revoked) = scoped_revoke_external_workdir_grant( - State(api.clone()), - Extension(test_browser_request_actor()), - AxumPath((TEST_WORKSPACE_ID.to_string(), grant.grant_id.clone())), + let revoker = { + let api = api.clone(); + let grant_id = grant.grant_id.clone(); + tokio::spawn(async move { + scoped_revoke_external_workdir_grant( + State(api), + Extension(test_browser_request_actor()), + AxumPath((TEST_WORKSPACE_ID.to_string(), grant_id)), + ) + .await + }) + }; + let (operation_id, operation) = loop { + let text = second.next().await.unwrap().unwrap().into_text().unwrap(); + let frame = serde_json::from_str::(&text).unwrap(); + if let ExternalWorkdirServerMessage::Operation { + generation: 2, + operation_id, + operation, + } = frame.message + { + break (operation_id, operation); + } + }; + assert!(matches!( + operation.as_inner(), + WorkdirSessionOperation::CommandCancel(_) + )); + let result = workdir::http::dispatch_workdir_session_operation( + &local_session, + operation.into_inner(), + ) + .await + .unwrap() + .try_into() + .unwrap(); + second + .send(Message::Text( + serde_json::to_string(&ExternalWorkdirProviderFrame::current( + ExternalWorkdirProviderMessage::OperationResult { + generation: 2, + operation_id, + outcome: ExternalWorkdirOperationOutcome::Completed { result }, + }, + )) + .unwrap() + .into(), + )) + .await + .unwrap(); + + let (operation_id, operation) = loop { + let text = second.next().await.unwrap().unwrap().into_text().unwrap(); + let frame = serde_json::from_str::(&text).unwrap(); + if let ExternalWorkdirServerMessage::Operation { + generation: 2, + operation_id, + operation, + } = frame.message + { + break (operation_id, operation); + } + }; + assert!(matches!( + operation.as_inner(), + WorkdirSessionOperation::CommandOutput(workdir::CommandOutputRequest { + wait: true, + .. + }) + )); + let provider_result = workdir::http::dispatch_workdir_session_operation( + &local_session, + operation.into_inner(), ) .await .unwrap(); + assert!(matches!( + provider_result, + WorkdirSessionOperationResult::CommandOutput(ref output) + if output.status != workdir::CommandStatus::Running + )); + let result = provider_result.try_into().unwrap(); + second + .send(Message::Text( + serde_json::to_string(&ExternalWorkdirProviderFrame::current( + ExternalWorkdirProviderMessage::OperationResult { + generation: 2, + operation_id, + outcome: ExternalWorkdirOperationOutcome::Completed { result }, + }, + )) + .unwrap() + .into(), + )) + .await + .unwrap(); + let Json(revoked) = revoker.await.unwrap().unwrap(); assert_eq!(revoked.status, "revoked"); let active_links = api .store @@ -27545,6 +28087,16 @@ mod tests { revoke.message, ExternalWorkdirServerMessage::Revoke { generation: 2, .. } )); + second + .send(Message::Text( + serde_json::to_string(&ExternalWorkdirProviderFrame::current( + ExternalWorkdirProviderMessage::RevokeAcknowledged { generation: 2 }, + )) + .unwrap() + .into(), + )) + .await + .unwrap(); server.abort(); } @@ -27559,6 +28111,7 @@ mod tests { expires_at: Some(Utc::now() + Duration::minutes(1)), capabilities: workdir::WorkdirSessionCapabilities::READ_ONLY, admission: Arc::new(tokio::sync::Semaphore::new(16)), + shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)), sender, }); let (sender, _) = tokio::sync::mpsc::channel(1); diff --git a/crates/yoi/src/workdir_share.rs b/crates/yoi/src/workdir_share.rs index 82d74ec8b..196863b3f 100644 --- a/crates/yoi/src/workdir_share.rs +++ b/crates/yoi/src/workdir_share.rs @@ -14,9 +14,10 @@ use server_api::{ use tokio_tungstenite::connect_async_with_config; use tokio_tungstenite::tungstenite::{Message, protocol::WebSocketConfig}; use workdir::external::{ - ExternalProviderInstanceId, ExternalWorkdirGrantId, ExternalWorkdirOperationId, - ExternalWorkdirOperationOutcome, ExternalWorkdirProviderFrame, ExternalWorkdirProviderMessage, - ExternalWorkdirProviderRegistration, ExternalWorkdirServerFrame, ExternalWorkdirServerMessage, + ExternalProviderInstanceId, ExternalWorkdirGrantId, ExternalWorkdirOperationError, + ExternalWorkdirOperationId, ExternalWorkdirOperationOutcome, ExternalWorkdirProviderFrame, + ExternalWorkdirProviderMessage, ExternalWorkdirProviderRegistration, + ExternalWorkdirServerFrame, ExternalWorkdirServerMessage, }; use workdir::http::{WorkdirTransportError, dispatch_workdir_session_operation}; use workdir::{ @@ -796,10 +797,12 @@ async fn serve_provider_connection( }), ) .unwrap_or_else(|_| ExternalWorkdirOperationOutcome::Failed { - error: WorkdirTransportError { - code: workdir::http::WorkdirTransportErrorCode::Internal, - message: "External Workdir operation executor failed".to_string(), - }, + error: ExternalWorkdirOperationError::from_transport_error( + WorkdirTransportError { + code: workdir::http::WorkdirTransportErrorCode::Internal, + message: "External Workdir operation executor failed".to_string(), + }, + ), }); let _ = completion_sender.blocking_send(OperationCompletion { operation_id: task_operation_id, @@ -912,14 +915,16 @@ async fn execute_operation( Ok(result) => match result.try_into() { Ok(result) => ExternalWorkdirOperationOutcome::Completed { result }, Err(message) => ExternalWorkdirOperationOutcome::Failed { - error: WorkdirTransportError { + error: ExternalWorkdirOperationError::from_transport_error(WorkdirTransportError { code: workdir::http::WorkdirTransportErrorCode::Unsupported, message, - }, + }), }, }, Err(error) => ExternalWorkdirOperationOutcome::Failed { - error: WorkdirTransportError::from_workdir_error(&error), + error: ExternalWorkdirOperationError::from_transport_error( + WorkdirTransportError::from_workdir_error(&error), + ), }, } } From 1980da288009701a50e39ff41a9536af4d8615fb Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 12:37:13 +0900 Subject: [PATCH 05/11] fix: close early dispatch recovery gaps --- crates/agen/src/engine.rs | 332 +++++++++++-- .../llm_client/scheme/anthropic/request.rs | 1 + .../src/llm_client/scheme/gemini/request.rs | 1 + .../llm_client/scheme/openai_chat/request.rs | 1 + .../scheme/openai_responses/request.rs | 1 + crates/agen/src/llm_client/types.rs | 25 +- crates/agen/src/timeline/timeline.rs | 85 +++- crates/agen/tests/annotated_history_test.rs | 1 + crates/agen/tests/parallel_execution_test.rs | 447 ++++++++++++++++++ crates/memory/src/extract/input.rs | 2 +- crates/session-store/src/logged_item.rs | 21 + crates/session-store/src/public_snapshot.rs | 4 +- crates/worker/src/session_capture.rs | 5 +- crates/worker/src/worker.rs | 12 +- 14 files changed, 868 insertions(+), 70 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index fa4ab25a6..8dd5825a3 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -98,31 +98,39 @@ pub struct EngineConfig { /// Project terminal tool outputs into the assistant's original ToolCall order. /// /// Runtime history intentionally retains completion order so every result can -/// be committed without waiting for slower siblings. The provider projection -/// is deterministic within each contiguous result batch and does not rewrite -/// the committed transcript. +/// be committed without waiting for slower siblings. An internal durable +/// response boundary lets this projection move every early result after the +/// complete assistant response, even when trailing text/reasoning arrived after +/// a fast result. Boundaries are never exposed to providers. struct ProviderHistoryProjection { items: Vec, original_to_projected_index: Vec, } -fn materialize_provider_history(items: &[Item]) -> ProviderHistoryProjection { - let mut materialized: Vec<_> = items.iter().cloned().enumerate().collect(); - let mut call_order = HashMap::::new(); - let mut next_call_order = 0usize; - - for (_, item) in &materialized { - if let Item::ToolCall { +fn normalize_tool_run(run: &mut [(usize, Item)], call_order: &HashMap) { + run.sort_by_key(|(_, item)| match item { + Item::ToolCall { call_id, call_index, .. - } = item - { - call_order.insert(call_id.clone(), call_index.unwrap_or(next_call_order)); - next_call_order += 1; - } - } + } => ( + 0usize, + call_index + .or_else(|| call_order.get(call_id).copied()) + .unwrap_or(usize::MAX), + ), + Item::ToolResult { call_id, .. } => ( + 1usize, + call_order.get(call_id).copied().unwrap_or(usize::MAX), + ), + _ => unreachable!("tool run contains only calls and results"), + }); +} +fn normalize_legacy_history( + materialized: &mut [(usize, Item)], + call_order: &HashMap, +) { let mut index = 0usize; while index < materialized.len() { if !matches!( @@ -141,29 +149,128 @@ fn materialize_provider_history(items: &[Item]) -> ProviderHistoryProjection { { index += 1; } - materialized[start..index].sort_by_key(|(_, item)| match item { - Item::ToolCall { - call_id, - call_index, - .. - } => ( - 0usize, - call_index - .or_else(|| call_order.get(call_id).copied()) - .unwrap_or(usize::MAX), - ), - Item::ToolResult { call_id, .. } => ( - 1usize, - call_order.get(call_id).copied().unwrap_or(usize::MAX), - ), - _ => unreachable!("tool run contains only calls and results"), - }); + normalize_tool_run(&mut materialized[start..index], call_order); + } +} + +fn normalize_response_history(response: &mut [(usize, Item)], call_order: &HashMap) { + // Provider-confirmed call completions can be observed out of index order. + // Reorder only the values occupying ToolCall slots so surrounding assistant + // text/reasoning retains its streamed order. + let mut calls = response + .iter() + .filter_map(|(_, item)| match item { + Item::ToolCall { .. } => Some(item.clone()), + _ => None, + }) + .collect::>(); + calls.sort_by_key(|item| match item { + Item::ToolCall { + call_id, + call_index, + .. + } => call_index + .or_else(|| call_order.get(call_id).copied()) + .unwrap_or(usize::MAX), + _ => unreachable!("filtered ToolCall list"), + }); + let mut calls = calls.into_iter(); + for (_, item) in response.iter_mut() { + if matches!(item, Item::ToolCall { .. }) { + *item = calls.next().expect("one replacement per ToolCall slot"); + } } - let mut original_to_projected_index = vec![0; materialized.len()]; + // A host-injected user/system item is a hard provider turn boundary. Do not + // move a result across it; normalize each assistant-side region separately. + let mut start = 0usize; + for index in 0..=response.len() { + let at_boundary = index == response.len() + || matches!( + response[index].1, + Item::Message { + role: crate::llm_client::types::Role::User + | crate::llm_client::types::Role::System, + .. + } + ); + if at_boundary { + response[start..index].sort_by_key(|(_, item)| match item { + Item::ToolResult { call_id, .. } => ( + 1usize, + call_order.get(call_id).copied().unwrap_or(usize::MAX), + ), + _ => (0usize, 0usize), + }); + start = index.saturating_add(1); + } + } +} + +fn materialize_provider_history(items: &[Item]) -> ProviderHistoryProjection { + let mut call_order = HashMap::::new(); + let mut next_call_order = 0usize; + for item in items { + if let Item::ToolCall { + call_id, + call_index, + .. + } = item + { + call_order.insert(call_id.clone(), call_index.unwrap_or(next_call_order)); + next_call_order += 1; + } + } + + let first_boundary = items + .iter() + .position(|item| matches!(item, Item::AssistantResponseBoundary { .. })); + let legacy_end = first_boundary.unwrap_or(items.len()); + let mut legacy = items[..legacy_end] + .iter() + .cloned() + .enumerate() + .collect::>(); + normalize_legacy_history(&mut legacy, &call_order); + + let mut materialized = legacy; + let mut boundary_indices = Vec::new(); + let mut index = legacy_end; + while index < items.len() { + debug_assert!(matches!( + items[index], + Item::AssistantResponseBoundary { .. } + )); + boundary_indices.push(index); + let start = index + 1; + let end = items[start..] + .iter() + .position(|item| matches!(item, Item::AssistantResponseBoundary { .. })) + .map(|offset| start + offset) + .unwrap_or(items.len()); + let mut response = items[start..end] + .iter() + .cloned() + .enumerate() + .map(|(offset, item)| (start + offset, item)) + .collect::>(); + normalize_response_history(&mut response, &call_order); + materialized.extend(response); + index = end; + } + + let mut original_to_projected_index = vec![0; items.len()]; for (projected_index, (original_index, _)) in materialized.iter().enumerate() { original_to_projected_index[*original_index] = projected_index; } + for boundary_index in boundary_indices { + let preceding_projected = materialized + .iter() + .take_while(|(original_index, _)| *original_index < boundary_index) + .count() + .saturating_sub(1); + original_to_projected_index[boundary_index] = preceding_projected; + } ProviderHistoryProjection { items: materialized.into_iter().map(|(_, item)| item).collect(), @@ -253,6 +360,7 @@ type ToolCallExecutionInfo = ( /// provider-confirmed blocks complete. struct EarlyToolExecutionBatch { batch_id: String, + turn_id: usize, calls: Vec<(usize, ToolCall)>, seen_call_ids: HashSet, call_info_map: HashMap, @@ -266,9 +374,10 @@ struct EarlyToolExecutionBatch { } impl EarlyToolExecutionBatch { - fn new(batch_id: String) -> Self { + fn new(batch_id: String, turn_id: usize) -> Self { Self { batch_id, + turn_id, calls: Vec::new(), seen_call_ids: HashSet::new(), call_info_map: HashMap::new(), @@ -1362,7 +1471,13 @@ impl Engine { let context = ToolExecutionContext::new(&tool_call.id, &batch.batch_id, call_index); if !tool_call.input.is_object() { - self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + Some(context.execution_id()), + )?; batch.synthetic_results.push_back(ToolResult::error( &tool_call.id, "Tool arguments must be one complete JSON object", @@ -1372,7 +1487,13 @@ impl Engine { } let Some((meta, tool)) = self.tool_server.get_tool(&tool_call.name) else { - self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + Some(context.execution_id()), + )?; batch.synthetic_results.push_back(ToolResult::error( &tool_call.id, format!("Tool not found: {}", tool_call.name), @@ -1502,12 +1623,24 @@ impl Engine { ) .into()); } - self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + Some(context.execution_id()), + )?; batch.synthetic_results.push_back(result); batch.dispatched_any = true; } PreToolAction::Skip => { - self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + Some(context.execution_id()), + )?; batch.synthetic_results.push_back(ToolResult::error( &tool_call.id, "Tool execution skipped by interceptor", @@ -1515,7 +1648,13 @@ impl Engine { batch.dispatched_any = true; } PreToolAction::Abort(reason) => { - self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + Some(context.execution_id()), + )?; batch.synthetic_results.push_back(ToolResult::error( &tool_call.id, format!("Tool execution aborted: {reason}"), @@ -1544,6 +1683,8 @@ impl Engine { annotate, result, None, + batch.turn_id, + true, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1569,6 +1710,8 @@ impl Engine { annotate, result, None, + batch.turn_id, + false, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1615,6 +1758,8 @@ impl Engine { annotate, result, Some(&attempt_id), + batch.turn_id, + false, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1636,6 +1781,8 @@ impl Engine { annotate, ToolResult::outcome_unknown(call_id), Some(attempt_id), + batch.turn_id, + false, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1665,6 +1812,8 @@ impl Engine { annotate, result, None, + batch.turn_id, + true, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1701,6 +1850,8 @@ impl Engine { annotate, result, Some(&attempt_id), + batch.turn_id, + true, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -1768,6 +1919,7 @@ impl Engine { // Map from tool call ID to (ToolCall, Meta, Tool, Context) // Retained because it's needed for PostToolCall hooks + let turn_id = self.turn_count.saturating_sub(1); let mut call_info_map = HashMap::new(); let mut synthetic_results = Vec::new(); let batch_id = format!("tool-batch-{}", self.tool_execution_batch_count); @@ -1952,6 +2104,8 @@ impl Engine { annotate, result, None, + turn_id, + true, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -1982,6 +2136,8 @@ impl Engine { annotate, result, Some(&attempt_id), + turn_id, + true, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -2061,6 +2217,8 @@ impl Engine { annotate, result, Some(&attempt_id), + turn_id, + false, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -2087,6 +2245,8 @@ impl Engine { annotate, ToolResult::outcome_unknown(call_id), Some(attempt_id), + turn_id, + false, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -2135,6 +2295,8 @@ impl Engine { annotate: &mut impl FnMut(&Item) -> Result, mut tool_result: ToolResult, execution_attempt_id: Option<&str>, + turn_id: usize, + invoke_post_hook: bool, call_info_map: &HashMap< String, ( @@ -2223,10 +2385,10 @@ impl Engine { ); self.emit_tool_result(&tool_result); - if let Some((tool_call, meta, tool, context)) = call_info { + if invoke_post_hook && let Some((tool_call, meta, tool, context)) = call_info { let invocation = self.interceptor_invocation( InterceptorPhase::PostToolCall, - Some(self.turn_count.saturating_sub(1)), + Some(turn_id), Some(InterceptorCallId::Tool(tool_call.id.clone())), context.call_index, ); @@ -2239,16 +2401,28 @@ impl Engine { tool: tool.clone(), context: context.clone(), }; - let post_tool_action = - self.interceptor - .post_tool_call(&info) - .await - .map_err(|error| { + let post_tool_action = tokio::select! { + result = self.interceptor.post_tool_call(&info) => { + result.map_err(|error| { EngineError::from(InterceptorFailure::new( InterceptorPhase::PostToolCall, error, )) - })?; + })? + } + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!(call_id = %tool_call.id, "Paused while awaiting post-tool policy"); + } + return Err(EngineError::PauseRequested); + } + cancel = self.cancel_rx.recv() => { + if cancel.is_some() { + info!(call_id = %tool_call.id, "Cancelled while awaiting post-tool policy"); + } + return Err(EngineError::Cancelled); + } + }; if let PostToolAction::Abort(reason) = post_tool_action { return Err(EngineError::Aborted(reason)); } @@ -2284,6 +2458,8 @@ impl Engine { annotate, ToolResult::outcome_unknown(&call_id), None, + self.turn_count, + false, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -2532,6 +2708,14 @@ impl Engine { "tool-call early dispatch deferred until response completion: interceptor requires the whole-response turn-end gate", ); } + if early_dispatch { + let response_id = format!("assistant-response-{current_turn}-{current_llm_call}"); + self.append_history_items( + history, + std::iter::once(Item::assistant_response_boundary(response_id)), + annotate, + )?; + } let assistant_start = history.len(); let StreamResponseOutput { completion: stream_outcome, @@ -2673,7 +2857,7 @@ impl Engine { if early_tools.is_none() { let batch_id = format!("tool-batch-{}", self.tool_execution_batch_count); self.tool_execution_batch_count += 1; - early_tools = Some(EarlyToolExecutionBatch::new(batch_id)); + early_tools = Some(EarlyToolExecutionBatch::new(batch_id, current_turn)); } let admission = self .admit_early_tool_calls( @@ -3057,7 +3241,7 @@ impl Engine { ); self.tool_execution_batch_count += 1; early_tools = - Some(EarlyToolExecutionBatch::new(batch_id)); + Some(EarlyToolExecutionBatch::new(batch_id, turn)); } let fresh_calls = { let batch = early_tools.as_ref().expect("batch initialized"); @@ -3211,6 +3395,8 @@ impl Engine { annotate, result, Some(&attempt_id), + batch.turn_id, + true, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, @@ -3941,6 +4127,7 @@ fn items_trace_payload( fn item_kind(item: &Item) -> &'static str { match item { Item::Message { .. } => "message", + Item::AssistantResponseBoundary { .. } => "assistant_response_boundary", Item::ToolCall { .. } => "tool_call", Item::ToolResult { .. } => "tool_result", Item::Reasoning { .. } => "reasoning", @@ -4025,6 +4212,51 @@ mod tests { assert_eq!(projection.original_to_projected_index, [1, 3, 0, 2]); } + #[test] + fn provider_projection_keeps_complete_response_before_early_results() { + let items = vec![ + Item::assistant_response_boundary("response-1"), + Item::tool_call_json("call_second", "second", serde_json::json!({})) + .with_tool_execution_metadata(1, Some("batch:call_second".to_string())) + .with_status(ItemStatus::InProgress), + Item::tool_result("call_second", "second result"), + Item::assistant_message("trailing text"), + Item::reasoning("trailing reasoning"), + Item::tool_call_json("call_first", "first", serde_json::json!({})) + .with_tool_execution_metadata(0, Some("batch:call_first".to_string())) + .with_status(ItemStatus::InProgress), + Item::tool_result("call_first", "first result"), + ]; + + let projection = materialize_provider_history(&items); + let order = projection + .items + .iter() + .map(|item| match item { + Item::Message { .. } => "text".to_string(), + Item::Reasoning { .. } => "reasoning".to_string(), + Item::ToolCall { call_id, .. } => format!("call:{call_id}"), + Item::ToolResult { call_id, .. } => format!("result:{call_id}"), + Item::AssistantResponseBoundary { .. } => panic!("boundary leaked"), + }) + .collect::>(); + assert_eq!( + order, + [ + "call:call_first", + "text", + "reasoning", + "call:call_second", + "result:call_first", + "result:call_second", + ] + ); + assert_eq!( + projection.original_to_projected_index, + [0, 0, 5, 1, 2, 3, 4] + ); + } + #[test] fn tool_attachment_round_trips_through_durable_history_json() { let body: Arc<[u8]> = Arc::from(&b"image-body"[..]); diff --git a/crates/agen/src/llm_client/scheme/anthropic/request.rs b/crates/agen/src/llm_client/scheme/anthropic/request.rs index 62fac8d17..b546f4c94 100644 --- a/crates/agen/src/llm_client/scheme/anthropic/request.rs +++ b/crates/agen/src/llm_client/scheme/anthropic/request.rs @@ -263,6 +263,7 @@ impl AnthropicScheme { for (i, item) in items.iter().enumerate() { match item { + Item::AssistantResponseBoundary { .. } => {} Item::Message { role, content, .. } => { let parts: Vec = content .iter() diff --git a/crates/agen/src/llm_client/scheme/gemini/request.rs b/crates/agen/src/llm_client/scheme/gemini/request.rs index e2fac9179..c1eb6d82c 100644 --- a/crates/agen/src/llm_client/scheme/gemini/request.rs +++ b/crates/agen/src/llm_client/scheme/gemini/request.rs @@ -242,6 +242,7 @@ impl GeminiScheme { for item in items { match item { + Item::AssistantResponseBoundary { .. } => {} Item::Message { role, content, .. } => { // Flush pending parts self.flush_pending_parts( diff --git a/crates/agen/src/llm_client/scheme/openai_chat/request.rs b/crates/agen/src/llm_client/scheme/openai_chat/request.rs index d971ba8b0..e8a87496e 100644 --- a/crates/agen/src/llm_client/scheme/openai_chat/request.rs +++ b/crates/agen/src/llm_client/scheme/openai_chat/request.rs @@ -221,6 +221,7 @@ impl OpenAIScheme { ); } match item { + Item::AssistantResponseBoundary { .. } => {} Item::Message { role, content, .. } => { // Flush pending tool calls self.flush_pending_assistant( diff --git a/crates/agen/src/llm_client/scheme/openai_responses/request.rs b/crates/agen/src/llm_client/scheme/openai_responses/request.rs index 63df50248..e26b60a51 100644 --- a/crates/agen/src/llm_client/scheme/openai_responses/request.rs +++ b/crates/agen/src/llm_client/scheme/openai_responses/request.rs @@ -254,6 +254,7 @@ fn convert_items_to_input(items: &[Item], supports_images: bool) -> Vec {} Item::Message { role, content, .. } => { let (role_str, text_variant): (&'static str, fn(String) -> InputContent) = match role { diff --git a/crates/agen/src/llm_client/types.rs b/crates/agen/src/llm_client/types.rs index 50ed7506a..1a9d3ca24 100644 --- a/crates/agen/src/llm_client/types.rs +++ b/crates/agen/src/llm_client/types.rs @@ -93,6 +93,12 @@ pub enum Item { status: Option, }, + /// Internal durable boundary used to reconstruct one assistant response when + /// early ToolResults are committed before the response stream has ended. + /// Provider request projections remove this item before serialization. + #[doc(hidden)] + AssistantResponseBoundary { response_id: String }, + /// Tool call from the assistant ToolCall { /// Optional item ID @@ -107,11 +113,12 @@ pub enum Item { /// Provider/model order among tool calls in this assistant response. #[serde(default, skip_serializing_if = "Option::is_none")] call_index: Option, - /// Durable identity of a tool execution attempt that has started. + /// Durable identity of an execution or synthetic terminalization that + /// owns this call. Once present, restore must never infer the call pending. #[serde(default, skip_serializing_if = "Option::is_none")] execution_id: Option, - /// Item status. `InProgress` on a persisted ToolCall means provider - /// execution ownership was transferred and must never be inferred pending. + /// Item status. `InProgress` on a persisted ToolCall means execution or + /// terminalization ownership was transferred and must never be inferred pending. #[serde(skip_serializing_if = "Option::is_none")] status: Option, }, @@ -226,6 +233,14 @@ impl Item { } } + /// Create an internal durable assistant-response boundary. + #[doc(hidden)] + pub fn assistant_response_boundary(response_id: impl Into) -> Self { + Self::AssistantResponseBoundary { + response_id: response_id.into(), + } + } + // ======================================================================== // Tool call constructors // ======================================================================== @@ -405,6 +420,7 @@ impl Item { pub fn with_id(mut self, id: impl Into) -> Self { match &mut self { Self::Message { id: item_id, .. } => *item_id = Some(id.into()), + Self::AssistantResponseBoundary { .. } => {} Self::ToolCall { id: item_id, .. } => *item_id = Some(id.into()), Self::ToolResult { id: item_id, .. } => *item_id = Some(id.into()), Self::Reasoning { id: item_id, .. } => *item_id = Some(id.into()), @@ -416,6 +432,7 @@ impl Item { pub fn with_status(mut self, new_status: ItemStatus) -> Self { match &mut self { Self::Message { status, .. } => *status = Some(new_status), + Self::AssistantResponseBoundary { .. } => {} Self::ToolCall { status, .. } => *status = Some(new_status), Self::ToolResult { .. } => {} // Result items don't have status Self::Reasoning { status, .. } => *status = Some(new_status), @@ -431,6 +448,7 @@ impl Item { pub fn id(&self) -> Option<&str> { match self { Self::Message { id, .. } => id.as_deref(), + Self::AssistantResponseBoundary { .. } => None, Self::ToolCall { id, .. } => id.as_deref(), Self::ToolResult { id, .. } => id.as_deref(), Self::Reasoning { id, .. } => id.as_deref(), @@ -441,6 +459,7 @@ impl Item { pub fn item_type(&self) -> &'static str { match self { Self::Message { .. } => "message", + Self::AssistantResponseBoundary { .. } => "assistant_response_boundary", Self::ToolCall { .. } => "tool_call", Self::ToolResult { .. } => "tool_result", Self::Reasoning { .. } => "reasoning", diff --git a/crates/agen/src/timeline/timeline.rs b/crates/agen/src/timeline/timeline.rs index fcdac1fc3..6dfe80b2c 100644 --- a/crates/agen/src/timeline/timeline.rs +++ b/crates/agen/src/timeline/timeline.rs @@ -374,7 +374,9 @@ pub struct Timeline { thinking_block_handlers: Vec>, tool_use_block_handlers: Vec>, - // 現在アクティブなブロック + // 現在アクティブなブロック。provider は index ごとに複数 scope を + // interleave できるため、単一の current block ではなく全 scope を追跡する。 + active_blocks: HashMap<(BlockType, usize), ()>, current_block: Option, // 1リクエスト内で受信した Usage event の集約バッファ。 @@ -400,6 +402,7 @@ impl Timeline { text_block_handlers: Vec::new(), thinking_block_handlers: Vec::new(), tool_use_block_handlers: Vec::new(), + active_blocks: HashMap::new(), current_block: None, pending_usage: None, } @@ -554,6 +557,8 @@ impl Timeline { } fn handle_block_start(&mut self, start: &BlockStart) { + self.active_blocks + .insert((start.block_type, start.index), ()); self.current_block = Some(start.block_type); let handlers = self.get_block_handlers_mut(start.block_type); @@ -567,10 +572,11 @@ impl Timeline { let block_type = delta.delta.block_type(); // OpenAIなどのプロバイダはBlockStartを送らない場合があるため、 - // Deltaが来たときにスコープがなければ暗黙的に開始する - if self.current_block.is_none() { - self.current_block = Some(block_type); - } + // Deltaが来たときにその index のスコープがなければ暗黙的に開始する + self.active_blocks + .entry((block_type, delta.index)) + .or_insert(()); + self.current_block = Some(block_type); let handlers = self.get_block_handlers_mut(block_type); for handler in handlers { @@ -588,7 +594,12 @@ impl Timeline { handler.dispatch_stop(stop); handler.end_scope(); } - self.current_block = None; + self.active_blocks.remove(&(stop.block_type, stop.index)); + self.current_block = self + .active_blocks + .keys() + .next() + .map(|(block_type, _)| *block_type); } fn handle_block_abort(&mut self, abort: &BlockAbort) { @@ -597,7 +608,12 @@ impl Timeline { handler.dispatch_abort(abort); handler.end_scope(); } - self.current_block = None; + self.active_blocks.remove(&(abort.block_type, abort.index)); + self.current_block = self + .active_blocks + .keys() + .next() + .map(|(block_type, _)| *block_type); } fn get_block_handlers_mut( @@ -617,19 +633,22 @@ impl Timeline { self.current_block } - /// 現在アクティブなブロックを中断する + /// 現在アクティブな全ブロックを中断する /// - /// キャンセルやエラー時に呼び出し、進行中のブロックに対して + /// キャンセルやエラー時に呼び出し、index ごとに進行中の全ブロックへ /// BlockAbortイベントを発火してスコープをクリーンアップする。 pub fn abort_current_block(&mut self) { - if let Some(block_type) = self.current_block { + let active = self.active_blocks.keys().copied().collect::>(); + for (block_type, index) in active { let abort = crate::timeline::event::BlockAbort { - index: 0, // インデックスは不明なので0 + index, block_type, reason: "Cancelled".to_string(), }; self.handle_block_abort(&abort); } + self.active_blocks.clear(); + self.current_block = None; } } @@ -674,6 +693,50 @@ mod tests { assert!(calls.lock().unwrap().is_empty()); } + #[test] + fn abort_clears_every_indexed_scope_before_indices_are_reused() { + #[derive(Default)] + struct Scope { + started: bool, + } + struct IndexedToolHandler { + reused_cleanly: Arc>>, + } + impl Handler for IndexedToolHandler { + type Scope = Scope; + + fn on_event(&mut self, scope: &mut Scope, event: &ToolUseBlockEvent) { + if matches!(event, ToolUseBlockEvent::Start(_)) { + self.reused_cleanly.lock().unwrap().push(!scope.started); + scope.started = true; + } + } + } + + let reused_cleanly = Arc::new(Mutex::new(Vec::new())); + let mut timeline = Timeline::new(); + timeline.on_tool_use_block(IndexedToolHandler { + reused_cleanly: reused_cleanly.clone(), + }); + + timeline.dispatch(&Event::tool_use_start(0, "first", "tool")); + timeline.dispatch(&Event::tool_use_start(1, "second", "tool")); + // Stopping one interleaved block must not hide its still-active sibling. + timeline.dispatch(&Event::tool_use_stop(0)); + assert_eq!(timeline.current_block(), Some(BlockType::ToolUse)); + timeline.abort_current_block(); + assert!(timeline.current_block().is_none()); + + // A later provider response may reuse the same content-block index. + timeline.dispatch(&Event::tool_use_start(1, "replacement", "tool")); + timeline.dispatch(&Event::tool_use_stop(1)); + + assert_eq!( + reused_cleanly.lock().unwrap().as_slice(), + [true, true, true] + ); + } + #[test] fn test_meta_event_dispatch() { // シンプルなテスト用構造体 diff --git a/crates/agen/tests/annotated_history_test.rs b/crates/agen/tests/annotated_history_test.rs index cf1be3c08..b7c391965 100644 --- a/crates/agen/tests/annotated_history_test.rs +++ b/crates/agen/tests/annotated_history_test.rs @@ -36,6 +36,7 @@ async fn run_preserves_item_annotations_without_projecting_them() { Role::Assistant => "assistant", Role::System => "system", }, + Item::AssistantResponseBoundary { .. } => "assistant_response_boundary", Item::ToolCall { .. } => "tool_call", Item::ToolResult { .. } => "tool_result", Item::Reasoning { .. } => "reasoning", diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index f6d445fb0..744e4111a 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -308,6 +308,68 @@ impl LlmClient for ControlledStreamClient { } } +#[derive(Clone)] +struct MultiControlledStreamClient { + receivers: Arc< + Mutex< + std::collections::VecDeque< + tokio::sync::mpsc::UnboundedReceiver>, + >, + >, + >, +} + +impl MultiControlledStreamClient { + fn new( + response_count: usize, + ) -> ( + Self, + Vec>>, + ) { + let mut senders = Vec::with_capacity(response_count); + let mut receivers = std::collections::VecDeque::with_capacity(response_count); + for _ in 0..response_count { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + senders.push(tx); + receivers.push_back(rx); + } + ( + Self { + receivers: Arc::new(Mutex::new(receivers)), + }, + senders, + ) + } +} + +#[async_trait] +impl LlmClient for MultiControlledStreamClient { + async fn stream(&self, _request: Request) -> Result { + let mut rx = + self.receivers + .lock() + .unwrap() + .pop_front() + .ok_or_else(|| ClientError::Api { + status: Some(500), + code: Some("multi_controlled_stream_exhausted".to_string()), + message: "controlled responses exhausted".to_string(), + retry_after: None, + })?; + Ok(Box::pin(futures::stream::poll_fn(move |cx| { + rx.poll_recv(cx) + }))) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + ToolCallCompletionSupport::PerBlock + } +} + #[derive(Clone)] struct BoundedStreamClient { receiver: Arc>>>>, @@ -576,6 +638,82 @@ async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_pause() { assert_eq!(probe.call_count(), 0); } +#[tokio::test] +async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_cancel() { + struct BlockingEarlyPolicy { + entered: Arc, + release: Arc, + } + + #[async_trait] + impl Interceptor for BlockingEarlyPolicy { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn pre_tool_call( + &self, + _info: &mut ToolCallInfo<'_>, + ) -> InterceptorResult { + self.entered.notify_one(); + self.release.notified().await; + Ok(PreToolAction::Continue) + } + } + + let (client, tx) = BoundedStreamClient::new(1); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let tool = SlowTool::new("cancelled_policy_tool", 0); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(BlockingEarlyPolicy { + entered: entered.clone(), + release, + }); + engine.register_tool(tool.definition()); + let cancel = engine.cancel_sender(); + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + + tx.send(Ok(Event::tool_use_start( + 0, + "call_cancelled_policy", + "cancelled_policy_tool", + ))) + .await + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))) + .await + .unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), entered.notified()) + .await + .expect("pre-tool policy entered"); + + tx.send(Ok(Event::ping())).await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), tx.send(Ok(Event::ping()))) + .await + .expect("background pump continues draining the bounded provider stream") + .unwrap(); + cancel.send(()).await.unwrap(); + drop(tx); + + let (output, _history) = tokio::time::timeout(Duration::from_secs(1), run) + .await + .expect("cancel does not deadlock the bounded stream pump") + .unwrap(); + assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) + )); + assert_eq!(probe.call_count(), 0); +} + #[tokio::test] async fn per_block_dispatch_starts_siblings_without_waiting_for_stream_or_prior_result() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); @@ -680,6 +818,106 @@ async fn early_terminal_result_commits_while_response_stream_is_open() { let _ = run.await.unwrap(); } +#[derive(Clone, Copy)] +enum BlockingPostControl { + Pause, + Cancel, +} + +async fn assert_blocking_post_hook_services_control(control: BlockingPostControl) { + struct BlockingPostPolicy { + entered: Arc, + release: Arc, + } + + #[async_trait] + impl Interceptor for BlockingPostPolicy { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn post_tool_call( + &self, + _info: &ToolResultInfo<'_>, + ) -> InterceptorResult { + self.entered.notify_one(); + self.release.notified().await; + Ok(PostToolAction::Continue) + } + } + + let (client, tx) = BoundedStreamClient::new(1); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let tool = SlowTool::new("blocking_post_tool", 0); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(BlockingPostPolicy { + entered: entered.clone(), + release, + }); + engine.register_tool(tool.definition()); + let pause = engine.pause_sender(); + let cancel = engine.cancel_sender(); + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + + tx.send(Ok(Event::tool_use_start( + 0, + "call_blocking_post", + "blocking_post_tool", + ))) + .await + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))) + .await + .unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), entered.notified()) + .await + .expect("post-tool policy entered after the terminal result was committed"); + + tx.send(Ok(Event::ping())).await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), tx.send(Ok(Event::ping()))) + .await + .expect("background pump continues draining while post-tool policy blocks") + .unwrap(); + match control { + BlockingPostControl::Pause => pause.send(()).await.unwrap(), + BlockingPostControl::Cancel => cancel.send(()).await.unwrap(), + } + drop(tx); + + let (output, history) = tokio::time::timeout(Duration::from_secs(1), run) + .await + .expect("control request interrupts the blocked post-tool policy") + .unwrap(); + match control { + BlockingPostControl::Pause => assert!(matches!(output.result, EngineRunExit::Paused)), + BlockingPostControl::Cancel => assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) + )), + } + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_blocking_post" + ))); +} + +#[tokio::test] +async fn blocking_post_tool_policy_services_pause_after_result_commit() { + assert_blocking_post_hook_services_control(BlockingPostControl::Pause).await; +} + +#[tokio::test] +async fn blocking_post_tool_policy_services_cancel_after_result_commit() { + assert_blocking_post_hook_services_control(BlockingPostControl::Cancel).await; +} + #[tokio::test] async fn reverse_stop_order_preserves_model_call_indexes() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); @@ -720,6 +958,119 @@ async fn reverse_stop_order_preserves_model_call_indexes() { assert_eq!(by_id.get("call_order_b"), Some(&1)); } +#[tokio::test] +async fn early_pre_and_post_hooks_share_the_exact_response_turn_identity() { + struct TurnRecordingPolicy { + events: tokio::sync::mpsc::UnboundedSender<(InterceptorPhase, String, u64)>, + } + + #[async_trait] + impl Interceptor for TurnRecordingPolicy { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_>, + ) -> InterceptorResult { + self.events + .send(( + InterceptorPhase::PreToolCall, + info.call.id.clone(), + info.invocation.turn_id.expect("pre hook turn").0, + )) + .unwrap(); + Ok(PreToolAction::Continue) + } + + async fn post_tool_call( + &self, + info: &ToolResultInfo<'_>, + ) -> InterceptorResult { + self.events + .send(( + InterceptorPhase::PostToolCall, + info.call.id.clone(), + info.invocation.turn_id.expect("post hook turn").0, + )) + .unwrap(); + Ok(PostToolAction::Continue) + } + } + + let (client, mut responses) = MultiControlledStreamClient::new(3); + let first = responses.remove(0); + let second = responses.remove(0); + let third = responses.remove(0); + let (event_tx, mut event_rx) = tokio::sync::mpsc::unbounded_channel(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(TurnRecordingPolicy { events: event_tx }); + engine.register_tool(SlowTool::new("turn_identity", 0).definition()); + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + + first + .send(Ok(Event::tool_use_start(0, "call_turn_0", "turn_identity"))) + .unwrap(); + first.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + first.send(Ok(Event::tool_use_stop(0))).unwrap(); + let pre_0 = event_rx.recv().await.unwrap(); + let post_0 = event_rx.recv().await.unwrap(); + assert_eq!( + pre_0, + (InterceptorPhase::PreToolCall, "call_turn_0".into(), 0) + ); + assert_eq!( + post_0, + (InterceptorPhase::PostToolCall, "call_turn_0".into(), 0) + ); + first + .send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(first); + + second + .send(Ok(Event::tool_use_start(0, "call_turn_1", "turn_identity"))) + .unwrap(); + second + .send(Ok(Event::tool_input_delta(0, r#"{}"#))) + .unwrap(); + second.send(Ok(Event::tool_use_stop(0))).unwrap(); + let pre_1 = event_rx.recv().await.unwrap(); + let post_1 = event_rx.recv().await.unwrap(); + assert_eq!( + pre_1, + (InterceptorPhase::PreToolCall, "call_turn_1".into(), 1) + ); + assert_eq!( + post_1, + (InterceptorPhase::PostToolCall, "call_turn_1".into(), 1) + ); + second + .send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(second); + + third.send(Ok(Event::text_delta(0, "done"))).unwrap(); + third + .send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(third); + let (output, _) = run.await.unwrap(); + assert!(matches!(output.result, EngineRunExit::Finished)); +} + #[tokio::test] async fn reverse_stop_order_projects_next_provider_request_in_model_order() { let first_response = vec![ @@ -885,6 +1236,102 @@ async fn persisted_started_call_is_closed_unknown_and_never_resumed() { ))); } +#[tokio::test] +async fn rejected_synthetic_result_commit_restores_as_non_runnable() { + struct SyntheticDenial; + + #[async_trait] + impl Interceptor for SyntheticDenial { + fn supports_early_tool_dispatch(&self) -> bool { + true + } + + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_>, + ) -> InterceptorResult { + Ok(PreToolAction::SyntheticResult(ToolResult::error( + &info.call.id, + "permission denied", + ))) + } + } + + let client = MockLlmClient::with_responses(vec![ + vec![ + Event::tool_use_start(0, "call_denied_restore", "denied_restore"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ], + vec![ + Event::text_delta(0, "restored safely"), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ], + ]) + .with_completion_support(ToolCallCompletionSupport::PerBlock); + let tool = SlowTool::new("denied_restore", 0); + let probe = tool.clone(); + let reject_once = Arc::new(AtomicUsize::new(1)); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.set_interceptor(SyntheticDenial); + engine.register_tool(tool.definition()); + let reject_result = reject_once.clone(); + engine.on_history_append(move |item| { + if matches!(item, Item::ToolResult { call_id, .. } if call_id == "call_denied_restore") + && reject_result.swap(0, Ordering::SeqCst) == 1 + { + return Err("injected synthetic result commit failure".to_string()); + } + Ok(()) + }); + let mut history = History::new(); + + let output = engine.run(&mut history, "run denied tool").await; + assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + EngineError::HistoryAppend(ref message) + )) if message == "injected synthetic result commit failure" + )); + assert_eq!(probe.call_count(), 0); + assert!(history.items().any(|item| matches!( + item, + Item::ToolCall { + call_id, + execution_id: Some(_), + status: Some(agen::llm_client::ItemStatus::InProgress), + .. + } if call_id == "call_denied_restore" + ))); + assert!(!history.items().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_denied_restore" + ))); + + let mut engine = output.engine; + let resumed = engine.resume(&mut history).await; + assert!(matches!(resumed, EngineRunExit::Finished)); + assert_eq!( + probe.call_count(), + 0, + "denied side effect must never run on restore" + ); + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call_denied_restore" + ))); +} + #[tokio::test] async fn abort_during_early_admission_commits_synthetic_terminal() { struct AbortAdmission; diff --git a/crates/memory/src/extract/input.rs b/crates/memory/src/extract/input.rs index 9cc2b0934..e2169df46 100644 --- a/crates/memory/src/extract/input.rs +++ b/crates/memory/src/extract/input.rs @@ -43,7 +43,7 @@ fn render_items(items: &[Item]) -> String { Item::ToolResult { summary, .. } => { lines.push(format!("[ToolResult] {summary}")); } - Item::Reasoning { .. } => {} + Item::AssistantResponseBoundary { .. } | Item::Reasoning { .. } => {} } } lines.join("\n\n") diff --git a/crates/session-store/src/logged_item.rs b/crates/session-store/src/logged_item.rs index 6855b0973..546110baa 100644 --- a/crates/session-store/src/logged_item.rs +++ b/crates/session-store/src/logged_item.rs @@ -50,6 +50,9 @@ pub enum LoggedItem { role: LoggedRole, content: Vec, }, + AssistantResponseBoundary { + response_id: String, + }, ToolCall { call_id: String, name: String, @@ -122,6 +125,9 @@ impl From<&Item> for LoggedItem { role: (*role).into(), content: content.iter().map(LoggedContentPart::from).collect(), }, + Item::AssistantResponseBoundary { response_id } => Self::AssistantResponseBoundary { + response_id: response_id.clone(), + }, Item::ToolCall { call_id, name, @@ -185,6 +191,9 @@ impl From for Item { content: content.into_iter().map(Into::into).collect(), status: None, }, + LoggedItem::AssistantResponseBoundary { response_id } => { + Item::assistant_response_boundary(response_id) + } LoggedItem::ToolCall { call_id, name, @@ -343,6 +352,18 @@ mod tests { } } + #[test] + fn assistant_response_boundary_round_trips_for_provider_projection() { + let original = Item::assistant_response_boundary("response-7"); + let logged: LoggedItem = (&original).into(); + let json = serde_json::to_string(&logged).unwrap(); + assert!(json.contains(r#""kind":"assistant_response_boundary""#)); + assert_eq!( + Item::from(serde_json::from_str::(&json).unwrap()), + original + ); + } + #[test] fn round_trip_tool_call() { let original = Item::tool_call("call_42", "get_weather", r#"{"city":"Tokyo"}"#) diff --git a/crates/session-store/src/public_snapshot.rs b/crates/session-store/src/public_snapshot.rs index 72ec0a140..9fae33a15 100644 --- a/crates/session-store/src/public_snapshot.rs +++ b/crates/session-store/src/public_snapshot.rs @@ -386,8 +386,8 @@ fn project_item(item: &LoggedItem) -> Option { }) .collect(), }), - // Hidden model reasoning is never observable. - LoggedItem::Reasoning { .. } => None, + // Internal response grouping and hidden model reasoning are never observable. + LoggedItem::AssistantResponseBoundary { .. } | LoggedItem::Reasoning { .. } => None, } } diff --git a/crates/worker/src/session_capture.rs b/crates/worker/src/session_capture.rs index 92f145185..01125c332 100644 --- a/crates/worker/src/session_capture.rs +++ b/crates/worker/src/session_capture.rs @@ -415,7 +415,7 @@ impl SessionCapture { search_text: text, }); } - Item::Reasoning { .. } => {} + Item::AssistantResponseBoundary { .. } | Item::Reasoning { .. } => {} } } @@ -713,6 +713,9 @@ fn render_item( ), } } + Item::AssistantResponseBoundary { .. } => { + format!("[{} Assistant response boundary omitted]", entry.id) + } Item::Reasoning { .. } => format!("[{} Reasoning omitted]", entry.id), }; truncate_chars(&text, max_bytes) diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 6ab90b0d0..666eddec1 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -1264,7 +1264,10 @@ fn active_run_checkpoint_entry( fn is_ai_materialized_item(item: &Item) -> bool { match item { Item::Message { role, .. } => *role == Role::Assistant, - Item::ToolCall { .. } | Item::ToolResult { .. } | Item::Reasoning { .. } => true, + Item::AssistantResponseBoundary { .. } + | Item::ToolCall { .. } + | Item::ToolResult { .. } + | Item::Reasoning { .. } => true, } } @@ -1302,7 +1305,11 @@ where let subject = worker_subject(annotation_writer.state.location().session_id); let origin = if item.is_tool_result() { WorkerHistoryProvenance::ToolOutput { worker: subject } - } else if item.is_assistant_message() || item.is_tool_call() || item.is_reasoning() { + } else if item.is_assistant_message() + || matches!(item, Item::AssistantResponseBoundary { .. }) + || item.is_tool_call() + || item.is_reasoning() + { WorkerHistoryProvenance::ModelOutput { worker: subject } } else { // Unknown user/system append paths fail closed. Trusted system @@ -7557,6 +7564,7 @@ fn write_overview_header(items: &[Item], out: &mut String) { for item in items { match item { Item::Message { .. } => messages += 1, + Item::AssistantResponseBoundary { .. } => {} Item::ToolCall { .. } => tool_calls += 1, Item::ToolResult { .. } => tool_results += 1, Item::Reasoning { .. } => reasoning += 1, From 4ab2a2b88aa3e7a039665d9897e3ceddcbb5aafa Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 13:02:16 +0900 Subject: [PATCH 06/11] style: format merged workspace server changes --- crates/workspace-server/src/server.rs | 21 ++++++--------------- 1 file changed, 6 insertions(+), 15 deletions(-) diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index ef2a1498d..3021be462 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -119,13 +119,13 @@ use workdir::workspace::{ WorkspaceWorkdirSessionOperationRequest, }; use workdir::{CommandHandle, WorkdirSessionHandle}; +#[cfg(test)] +use worker_runtime::catalog::WorkingDirectoryAttachmentStatus; use worker_runtime::http_server::{ RUNTIME_HTTP_PROTOCOL_MAX_VERSION, RUNTIME_HTTP_PROTOCOL_MIN_VERSION, RUNTIME_HTTP_PROTOCOL_VERSION, }; #[cfg(test)] -use worker_runtime::catalog::WorkingDirectoryAttachmentStatus; -#[cfg(test)] use worker_runtime::resource::BackendResourceError; use worker_runtime::resource::BackendResourceFetchRequest; use worker_runtime::worker_backend::{ProfileRuntimeWorkerFactory, WorkerRuntimeExecutionBackend}; @@ -22519,13 +22519,8 @@ async fn get_runtime_worker( .store .list_workdir_registry(&api.config.workspace_id, 500)?; let updated_at = record.updated_at.clone(); - let worker = project_worker_registry_projection( - &api, - Some(&worker), - &record, - links, - &workdirs, - )?; + let worker = + project_worker_registry_projection(&api, Some(&worker), &record, links, &workdirs)?; Ok(Json(server_api::RuntimeWorkerShowResponse { worker, updated_at, @@ -29729,8 +29724,7 @@ mod tests { .unwrap(), runtime_id: Some(EMBEDDED_WORKER_RUNTIME_ID.to_string()), resource_key: None, - availability: - protocol::subscription::SubscriptionWorkerAvailability::Observed, + availability: protocol::subscription::SubscriptionWorkerAvailability::Observed, subject_revision: 1, worker_state: None, state: protocol::subscription::SubscriptionWorkerState::Paused, @@ -43454,10 +43448,7 @@ mod tests { selector: protocol::subscription::EventSubscriptionSelector::WorkerProtocol { .. }, snapshot: - protocol::subscription::SubscriptionSnapshot::WorkerProtocol { - events, - .. - }, + protocol::subscription::SubscriptionSnapshot::WorkerProtocol { events, .. }, .. }, ) = frame.payload From 36295edc08fd6dc6263aebd73c24646541de4db2 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 13:43:36 +0900 Subject: [PATCH 07/11] fix: require terminal stream completion --- crates/agen/src/engine.rs | 42 ++++++- crates/agen/tests/parallel_execution_test.rs | 105 +++++++++++++++++- .../src/feature/builtin/flow_transition.rs | 8 +- 3 files changed, 151 insertions(+), 4 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 8dd5825a3..737fac9a1 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -1529,12 +1529,28 @@ impl Engine { if pause.is_some() { info!(call_id = %expected_tool_use_id, "Paused while awaiting pre-tool policy"); } + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + None, + )?; + batch.dispatched_any = true; return Ok(EarlyToolAdmission::Pause); } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!(call_id = %expected_tool_use_id, "Cancelled while awaiting pre-tool policy"); } + self.append_early_tool_call( + history, + annotate, + &tool_call, + call_index, + None, + )?; + batch.dispatched_any = true; return Err(EngineError::Cancelled); } }; @@ -1664,6 +1680,7 @@ impl Engine { } PreToolAction::Pause => { self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; + batch.dispatched_any = true; return Ok(EarlyToolAdmission::Pause); } } @@ -3192,6 +3209,7 @@ impl Engine { self.tool_call_collector.begin_response(); let mut stream = ResponseStreamPump::start(stream); let mut early_tools: Option = None; + let mut response_completed = false; let mut event_count: usize = 0; loop { @@ -3229,6 +3247,14 @@ impl Engine { json!({}), ); } + if matches!( + &event, + Event::Status(StatusEvent { + status: crate::llm_client::event::ResponseStatus::Completed, + }) + ) { + response_completed = true; + } self.emit_stream_event(turn, llm_call, &event); self.timeline.dispatch(&event); if early_dispatch { @@ -3375,7 +3401,21 @@ impl Engine { })); } } - None => break, + None if response_completed => break, + None => { + // A clean transport EOF is not an authoritative provider + // response boundary. Treat it exactly like an interrupted + // stream so any early side effect is terminalized and never + // followed by automatic continuation/regeneration. + self.timeline.flush_usage(); + return Ok(StreamResponseOutput { + completion: StreamCompletion::Interrupted { + reason: "LLM stream ended before provider completion status" + .to_string(), + }, + early_tools, + }); + } } } terminal = async { diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index 744e4111a..3638f48d5 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -260,6 +260,7 @@ impl Tool for ContextRecordingTool { struct ControlledStreamClient { receiver: Arc>>>>, completion_support: ToolCallCompletionSupport, + stream_calls: Arc, } impl ControlledStreamClient { @@ -274,15 +275,21 @@ impl ControlledStreamClient { Self { receiver: Arc::new(Mutex::new(Some(rx))), completion_support, + stream_calls: Arc::new(AtomicUsize::new(0)), }, tx, ) } + + fn stream_count(&self) -> usize { + self.stream_calls.load(Ordering::SeqCst) + } } #[async_trait] impl LlmClient for ControlledStreamClient { async fn stream(&self, _request: Request) -> Result { + self.stream_calls.fetch_add(1, Ordering::SeqCst); let mut rx = self .receiver .lock() @@ -633,9 +640,29 @@ async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_pause() { pause.send(()).await.unwrap(); drop(tx); - let (output, _history) = run.await.unwrap(); + let (output, mut history) = run.await.unwrap(); assert!(matches!(output.result, EngineRunExit::Paused)); assert_eq!(probe.call_count(), 0); + assert!(history.items().any(|item| matches!( + item, + Item::ToolCall { + call_id, + execution_id: None, + status: Some(agen::llm_client::ItemStatus::Completed), + .. + } if call_id == "call_blocked_policy" + ))); + + let mut restored = Engine::new(MockLlmClient::new(vec![Event::Status(StatusEvent { + status: ResponseStatus::Completed, + })])); + restored.register_tool(probe.definition()); + let _ = restored.resume(&mut history).await; + assert_eq!(probe.call_count(), 1); + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_blocked_policy" + ))); } #[tokio::test] @@ -703,7 +730,7 @@ async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_cancel() { cancel.send(()).await.unwrap(); drop(tx); - let (output, _history) = tokio::time::timeout(Duration::from_secs(1), run) + let (output, mut history) = tokio::time::timeout(Duration::from_secs(1), run) .await .expect("cancel does not deadlock the bounded stream pump") .unwrap(); @@ -712,6 +739,26 @@ async fn blocking_pre_tool_policy_keeps_receiving_stream_and_services_cancel() { EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) )); assert_eq!(probe.call_count(), 0); + assert!(history.items().any(|item| matches!( + item, + Item::ToolCall { + call_id, + execution_id: None, + status: Some(agen::llm_client::ItemStatus::Completed), + .. + } if call_id == "call_cancelled_policy" + ))); + + let mut restored = Engine::new(MockLlmClient::new(vec![Event::Status(StatusEvent { + status: ResponseStatus::Completed, + })])); + restored.register_tool(probe.definition()); + let _ = restored.resume(&mut history).await; + assert_eq!(probe.call_count(), 1); + assert!(history.items().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_cancelled_policy" + ))); } #[tokio::test] @@ -1539,6 +1586,60 @@ async fn interrupted_stream_after_early_start_terminalizes_without_continuation_ ); } +#[tokio::test] +async fn clean_eof_without_provider_completion_after_early_start_stops_without_continuation() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let client_probe = client.clone(); + let tool = BarrierTool::new("barrier_clean_eof"); + let probe = tool.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::OnToolCallComplete); + engine.register_tool(tool.definition()); + + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + tx.send(Ok(Event::tool_use_start( + 0, + "call_clean_eof", + "barrier_clean_eof", + ))) + .unwrap(); + tx.send(Ok(Event::tool_input_delta(0, r#"{}"#))).unwrap(); + tx.send(Ok(Event::tool_use_stop(0))).unwrap(); + tokio::time::timeout(Duration::from_secs(1), probe.started.notified()) + .await + .expect("tool starts before clean transport EOF"); + probe.release.notify_one(); + drop(tx); + + let (output, history) = run.await.unwrap(); + assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(EngineError::Client( + ClientError::Api { + code: Some(ref code), + .. + } + ))) if code == "early_tool_stream_interrupted" + )); + assert_eq!( + client_probe.stream_count(), + 1, + "must not open a continuation request" + ); + assert_eq!(probe.starts.load(Ordering::SeqCst), 1); + assert_eq!( + history + .items() + .filter(|item| matches!(item, Item::ToolResult { call_id, .. } if call_id == "call_clean_eof")) + .count(), + 1 + ); +} + #[tokio::test] async fn unsupported_provider_visibly_defers_opt_in_dispatch() { let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::ResponseComplete); diff --git a/crates/worker/src/feature/builtin/flow_transition.rs b/crates/worker/src/feature/builtin/flow_transition.rs index 163ac8928..b96d1919d 100644 --- a/crates/worker/src/feature/builtin/flow_transition.rs +++ b/crates/worker/src/feature/builtin/flow_transition.rs @@ -636,7 +636,7 @@ mod tests { use agen::llm_client::error::ClientError; use agen::llm_client::event::{ BlockDelta, BlockMetadata, BlockStart, BlockStop, BlockType, DeltaContent, - Event as LlmEvent, StopReason, + Event as LlmEvent, ResponseStatus, StatusEvent, StopReason, }; use agen::llm_client::types::Request; use agen::tool::ToolExecutionContext; @@ -906,6 +906,9 @@ mod tests { reasoning: None, stop_reason: Some(StopReason::ToolUse), })), + Ok(LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + })), ]))) } else { Ok(Box::pin(stream::iter(vec![ @@ -924,6 +927,9 @@ mod tests { reasoning: None, stop_reason: Some(StopReason::EndTurn), })), + Ok(LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + })), ]))) } } From d81c9a8154b01dfc5b6fc1b1f1487d3b3ebfd830 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 13:51:03 +0900 Subject: [PATCH 08/11] fix: authorize signed workers for Workdir removal --- crates/workspace-server/src/server.rs | 455 ++++++++++++++++++++++++-- 1 file changed, 433 insertions(+), 22 deletions(-) diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index e9ebf8a4a..cdb53613e 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -119,13 +119,13 @@ use workdir::workspace::{ WorkspaceWorkdirSessionOperationRequest, }; use workdir::{CommandHandle, WorkdirSessionHandle}; +#[cfg(test)] +use worker_runtime::catalog::WorkingDirectoryAttachmentStatus; use worker_runtime::http_server::{ RUNTIME_HTTP_PROTOCOL_MAX_VERSION, RUNTIME_HTTP_PROTOCOL_MIN_VERSION, RUNTIME_HTTP_PROTOCOL_VERSION, }; #[cfg(test)] -use worker_runtime::catalog::WorkingDirectoryAttachmentStatus; -#[cfg(test)] use worker_runtime::resource::BackendResourceError; use worker_runtime::resource::BackendResourceFetchRequest; use worker_runtime::worker_backend::{ProfileRuntimeWorkerFactory, WorkerRuntimeExecutionBackend}; @@ -16765,13 +16765,19 @@ fn workdir_removal_source_actor( fn contract_workdir_removal_source_actor( context: &server_api::ServerRequestContext, ) -> ApiResult { - workdir_removal_source_actor( - context - .worker_source - .as_ref() - .map(|source| (source.runtime_id.as_str(), source.worker_id.as_str())), - context.actor.as_ref(), - ) + let worker_source = context + .worker_source + .as_ref() + .map(|source| (source.runtime_id.as_str(), source.worker_id.as_str())) + .or_else(|| { + context.runtime_source.as_ref().and_then(|source| { + source + .worker_id + .as_deref() + .map(|worker_id| (source.runtime_id.as_str(), worker_id)) + }) + }); + workdir_removal_source_actor(worker_source, context.actor.as_ref()) } async fn cleanup_runtime_working_directory_contract( @@ -22207,13 +22213,8 @@ async fn get_runtime_worker( .store .list_workdir_registry(&api.config.workspace_id, 500)?; let updated_at = record.updated_at.clone(); - let worker = project_worker_registry_projection( - &api, - Some(&worker), - &record, - links, - &workdirs, - )?; + let worker = + project_worker_registry_projection(&api, Some(&worker), &record, links, &workdirs)?; Ok(Json(server_api::RuntimeWorkerShowResponse { worker, updated_at, @@ -28232,6 +28233,36 @@ mod tests { .unwrap(); } + fn runtime_source_request( + identity: &worker_runtime::auth::RuntimeIdentityMaterial, + worker_id: Option<&str>, + method: &str, + path: &str, + body: Vec, + ) -> Request { + let proof = worker_runtime::auth::RuntimeRequestSourceSigner::from_identity(identity) + .issue( + "server-test", + TEST_WORKSPACE_ID, + worker_id, + worker_runtime::auth::WORKSPACE_REQUEST_PERMISSION, + method, + path, + &body, + i64::try_from(worker_runtime::auth::unix_now_seconds()).unwrap_or(i64::MAX), + 30, + ) + .unwrap(); + let mut request = Request::builder().method(method).uri(path).header( + worker_runtime::auth::RUNTIME_REQUEST_SOURCE_PROOF_HEADER, + proof, + ); + if !body.is_empty() { + request = request.header(CONTENT_TYPE, "application/json"); + } + request.body(Body::from(body)).unwrap() + } + fn runtime_resource_fetch_request( api: &WorkspaceApi, identity: &worker_runtime::auth::RuntimeIdentityMaterial, @@ -28733,8 +28764,7 @@ mod tests { .unwrap(), runtime_id: Some(EMBEDDED_WORKER_RUNTIME_ID.to_string()), resource_key: None, - availability: - protocol::subscription::SubscriptionWorkerAvailability::Observed, + availability: protocol::subscription::SubscriptionWorkerAvailability::Observed, subject_revision: 1, worker_state: None, state: protocol::subscription::SubscriptionWorkerState::Paused, @@ -37673,6 +37703,390 @@ mod tests { ); } + #[test] + fn workdir_removal_source_actor_preserves_verified_worker_and_account_sources() { + let actor = RequestActor { + user_id: "account-user".to_string(), + account_id: "account-test".to_string(), + handle: "account".to_string(), + display_name: "Account".to_string(), + auth_method: ActorAuthMethod::ApiToken, + }; + let context = |worker_source, runtime_source, actor| server_api::ServerRequestContext { + actor, + worker_source, + runtime_source, + origin: None, + transport_headers: Vec::new(), + }; + + assert_eq!( + contract_workdir_removal_source_actor(&context( + Some(server_api::ServerWorkerSource { + runtime_id: "legacy-runtime".to_string(), + worker_id: "legacy-worker".to_string(), + }), + Some(server_api::ServerRuntimeSource { + runtime_id: "runtime-source".to_string(), + worker_id: Some("runtime-worker".to_string()), + }), + Some(actor.clone()), + )) + .unwrap(), + "worker:legacy-runtime:legacy-worker" + ); + assert_eq!( + contract_workdir_removal_source_actor(&context( + None, + Some(server_api::ServerRuntimeSource { + runtime_id: "runtime-source".to_string(), + worker_id: Some("runtime-worker".to_string()), + }), + Some(actor.clone()), + )) + .unwrap(), + "worker:runtime-source:runtime-worker" + ); + assert_eq!( + contract_workdir_removal_source_actor(&context( + None, + Some(server_api::ServerRuntimeSource { + runtime_id: "runtime-only".to_string(), + worker_id: None, + }), + Some(actor.clone()), + )) + .unwrap(), + "account:account-test" + ); + let error = contract_workdir_removal_source_actor(&context( + None, + Some(server_api::ServerRuntimeSource { + runtime_id: "runtime-only".to_string(), + worker_id: None, + }), + None, + )) + .unwrap_err(); + assert_eq!(error.into_response().status(), StatusCode::FORBIDDEN); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn signed_worker_workdir_removal_reaches_both_generated_routes_and_records_actor() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let mut api = test_api(workspace.path()).await; + let identity = + worker_runtime::auth::RuntimeIdentityMaterial::generate("runtime-source").unwrap(); + configure_runtime_request_auth(&mut api, &identity, "runtime-source"); + seed_worker_source_member(&api, "runtime-source", "worker-source"); + let provider_runtime_id = "runtime-provider"; + let fixtures = [ + ( + "workspace-route-workdir", + format!("/api/w/{TEST_WORKSPACE_ID}/working-directories/workspace-route-workdir"), + "remove through workspace route", + ), + ( + "runtime-route-workdir", + format!( + "/api/w/{TEST_WORKSPACE_ID}/runtimes/{provider_runtime_id}/working-directories/runtime-route-workdir" + ), + "remove through Runtime route", + ), + ]; + let runtime = worker_runtime::Runtime::with_execution_backend( + worker_runtime::RuntimeOptions::default(), + Arc::new(DeterministicExecutionBackend::default()), + ) + .unwrap(); + let repository_source = server_api::RepositorySource { + kind: server_api::RepositorySourceKind::LocalPath, + uri: workspace.path().display().to_string(), + }; + let source_fingerprint = repository_source_fingerprint(&repository_source); + for (workdir_id, _, _) in &fixtures { + runtime + .create_working_directory(WorkingDirectoryRequest { + repository: WorkingDirectoryRepository { + id: "repo-test".to_string(), + provider: "git".to_string(), + source: repository_source.clone(), + source_revision: 1, + source_fingerprint: source_fingerprint.clone(), + selector: Some(RuntimeRepositorySelector("HEAD".to_string())), + }, + display_name: None, + materializer: Default::default(), + backend_workdir_id: Some((*workdir_id).to_string()), + materialization: None, + }) + .unwrap(); + seed_cleanup_workdir_for_runtime( + &api, + workdir_id, + provider_runtime_id, + "present", + "clean", + ); + } + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let provider_base_url = format!("http://{}", listener.local_addr().unwrap()); + let provider_server = tokio::spawn(serve_runtime_http_with_injected_test_auth( + runtime, listener, + )); + let now = now_registry_timestamp(); + let provider_identity = RuntimeIdentityMaterial::generate(provider_runtime_id).unwrap(); + let provider_binding = WorkspaceRuntimeBinding { + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: provider_runtime_id.to_string(), + display_name: "Runtime Provider".to_string(), + base_url: provider_base_url.clone(), + public_key: provider_identity.public_key, + public_key_fingerprint: String::new(), + binding_revision: 1, + state: StoredRuntimeBindingState::Verified, + authentication_mode: StoredRuntimeAuthenticationMode::LegacyServerIssuer, + workspace_key_id: None, + workspace_key_generation: None, + created_at: now.clone(), + updated_at: now, + revoked_at: None, + }; + api.store + .upsert_workspace_runtime_binding_record(provider_binding.clone(), false) + .await + .unwrap(); + api.runtime.register_or_replace( + RemoteWorkerRuntime::new( + RemoteRuntimeConfig { + runtime_id: provider_runtime_id.to_string(), + workspace_id: Some(TEST_WORKSPACE_ID.to_string()), + display_name: "Runtime Provider".to_string(), + base_url: provider_base_url, + bearer_token: Some(TEST_RUNTIME_HTTP_TOKEN.to_string()), + workspace_authorization: None, + strict_public_egress: false, + cached_worker_creation_available: true, + cached_os: "linux".to_string(), + cached_arch: "x86_64".to_string(), + cached_status: "active".to_string(), + timeout: std::time::Duration::from_secs(2), + }, + TEST_WORKSPACE_ID.to_string(), + "http://127.0.0.1:1".to_string(), + ) + .unwrap(), + ); + let observed = api + .runtime + .working_directory(provider_runtime_id, fixtures[0].0) + .unwrap(); + assert_eq!( + observed.state, + InternalWorkerOperationState::Accepted, + "{observed:?}" + ); + let app = build_router(api.clone()); + + for (workdir_id, path, reason) in fixtures { + let body = serde_json::to_vec(&json!({ "reason": reason })).unwrap(); + let response = app + .clone() + .oneshot(runtime_source_request( + &identity, + Some("worker-source"), + "DELETE", + &path, + body, + )) + .await + .unwrap(); + let status = response.status(); + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + assert_eq!( + status, + StatusCode::OK, + "{path}: {}", + String::from_utf8_lossy(&bytes) + ); + let response: WorkingDirectoryRemovalResponse = serde_json::from_slice(&bytes).unwrap(); + assert_eq!( + response.disposition, + WorkingDirectoryRemovalDisposition::Removed, + "{response:?}" + ); + assert!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, workdir_id) + .unwrap() + .is_none() + ); + let operation = api + .config_store + .find_workdir_removal_operation_by_intent( + TEST_WORKSPACE_ID, + workdir_id, + "worker:runtime-source:worker-source", + reason, + ) + .unwrap() + .unwrap(); + assert_eq!( + operation.source_actor, + "worker:runtime-source:worker-source" + ); + assert_eq!(operation.state, WorkdirRemovalOperationState::Completed); + assert_eq!( + operation.disposition, + Some(WorkdirRemovalDisposition::Removed) + ); + } + provider_server.abort(); + } + + #[tokio::test] + async fn workdir_removal_rejects_unverified_or_runtime_only_sources_before_removal() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let mut api = test_api(workspace.path()).await; + let identity = + worker_runtime::auth::RuntimeIdentityMaterial::generate("runtime-source").unwrap(); + configure_runtime_request_auth(&mut api, &identity, "runtime-source"); + let fixtures = [ + ("runtime-only-workdir", "runtime only"), + ("invalid-proof-workdir", "invalid proof"), + ("unauthenticated-workdir", "unauthenticated"), + ]; + for (workdir_id, _) in fixtures { + seed_cleanup_workdir_for_runtime( + &api, + workdir_id, + EMBEDDED_WORKER_RUNTIME_ID, + "present", + "clean", + ); + } + let app = build_router(api.clone()); + + let runtime_only_path = + format!("/api/w/{TEST_WORKSPACE_ID}/working-directories/runtime-only-workdir"); + let runtime_only = app + .clone() + .oneshot(runtime_source_request( + &identity, + None, + "DELETE", + &runtime_only_path, + serde_json::to_vec(&json!({ "reason": "runtime only" })).unwrap(), + )) + .await + .unwrap(); + assert_eq!(runtime_only.status(), StatusCode::FORBIDDEN); + + let invalid_path = + format!("/api/w/{TEST_WORKSPACE_ID}/working-directories/invalid-proof-workdir"); + let invalid = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri(&invalid_path) + .header(CONTENT_TYPE, "application/json") + .header( + worker_runtime::auth::RUNTIME_REQUEST_SOURCE_PROOF_HEADER, + "not-a-valid-proof", + ) + .body(Body::from(json!({ "reason": "invalid proof" }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(invalid.status(), StatusCode::UNAUTHORIZED); + + let unauthenticated_path = + format!("/api/w/{TEST_WORKSPACE_ID}/working-directories/unauthenticated-workdir"); + let unauthenticated = app + .oneshot( + Request::builder() + .method("DELETE") + .uri(&unauthenticated_path) + .header(CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ "reason": "unauthenticated" }).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(unauthenticated.status(), StatusCode::UNAUTHORIZED); + + for (workdir_id, _) in fixtures { + assert!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, workdir_id) + .unwrap() + .is_some() + ); + } + } + + #[tokio::test] + async fn account_workdir_removal_still_records_account_actor_and_preserves_guards() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + let token = seed_test_api_token(api.store.as_ref(), "workdir-removal"); + let workdir_id = "account-occupied-workdir"; + seed_cleanup_workdir_for_runtime( + &api, + workdir_id, + EMBEDDED_WORKER_RUNTIME_ID, + "present", + "clean", + ); + let occupant = seed_cleanup_worker(&api, 27, "pinned"); + seed_cleanup_link(&api, &occupant, workdir_id); + let reason = "retain occupied Workdir"; + let path = format!("/api/w/{TEST_WORKSPACE_ID}/working-directories/{workdir_id}"); + + let response = request_json_authenticated( + build_router(api.clone()), + "DELETE", + &path, + Some(json!({ "reason": reason })), + &token, + StatusCode::OK, + ) + .await; + assert_eq!(response["disposition"], "retained"); + assert!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, workdir_id) + .unwrap() + .is_some() + ); + let operation = api + .config_store + .find_workdir_removal_operation_by_intent( + TEST_WORKSPACE_ID, + workdir_id, + "account:account-workdir-removal", + reason, + ) + .unwrap() + .unwrap(); + assert_eq!(operation.source_actor, "account:account-workdir-removal"); + assert_eq!( + operation.disposition, + Some(WorkdirRemovalDisposition::Retained) + ); + assert_eq!( + operation.failure_category.as_deref(), + Some("blocked_by_live_authority") + ); + } + #[test] fn workdir_session_owner_is_only_sent_for_same_runtime_worker() { let embedded_worker = RuntimeWorkerRef::new("embedded-worker-runtime", "5"); @@ -42458,10 +42872,7 @@ mod tests { selector: protocol::subscription::EventSubscriptionSelector::WorkerProtocol { .. }, snapshot: - protocol::subscription::SubscriptionSnapshot::WorkerProtocol { - events, - .. - }, + protocol::subscription::SubscriptionSnapshot::WorkerProtocol { events, .. }, .. }, ) = frame.payload From 8ff7df7edfe4b998200ee8cbf3ba6e1c60511951 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 13:55:07 +0900 Subject: [PATCH 09/11] fix: preserve response-complete stream EOF --- crates/agen/src/engine.rs | 12 +- crates/agen/tests/parallel_execution_test.rs | 4 - .../response_completion_compatibility_test.rs | 142 ++++++++++++++++++ 3 files changed, 149 insertions(+), 9 deletions(-) create mode 100644 crates/agen/tests/response_completion_compatibility_test.rs diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 737fac9a1..45a28293d 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -3401,12 +3401,14 @@ impl Engine { })); } } - None if response_completed => break, + None if !early_dispatch || response_completed => break, None => { - // A clean transport EOF is not an authoritative provider - // response boundary. Treat it exactly like an interrupted - // stream so any early side effect is terminalized and never - // followed by automatic continuation/regeneration. + // Providers that support only response-complete tool calls use + // stream EOF as their established completion boundary. Early + // dispatch is never enabled for them. Once per-block dispatch is + // active, however, transport EOF alone cannot prove that the + // assistant response completed after a side effect may have begun. + // Require the provider's explicit terminal status in that mode. self.timeline.flush_usage(); return Ok(StreamResponseOutput { completion: StreamCompletion::Interrupted { diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index 3638f48d5..279c65eeb 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -1679,10 +1679,6 @@ async fn unsupported_provider_visibly_defers_opt_in_dispatch() { warning.contains("provider does not expose a trustworthy per-call completion boundary") })); - tx.send(Ok(Event::Status(StatusEvent { - status: ResponseStatus::Completed, - }))) - .unwrap(); drop(tx); tokio::time::timeout(Duration::from_secs(1), probe.started.notified()) .await diff --git a/crates/agen/tests/response_completion_compatibility_test.rs b/crates/agen/tests/response_completion_compatibility_test.rs new file mode 100644 index 000000000..c7b7d8fac --- /dev/null +++ b/crates/agen/tests/response_completion_compatibility_test.rs @@ -0,0 +1,142 @@ +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use agen::llm_client::{ + ClientError, LlmClient, Request, ResponseStream, ToolCallCompletionSupport, + scheme::{Scheme, gemini::GeminiScheme, openai_chat::OpenAIScheme}, +}; +use agen::{Engine, EngineRunExit, History, ToolCallDispatchMode}; +use async_trait::async_trait; + +#[derive(Clone)] +struct SseFrame { + event_type: &'static str, + data: &'static str, +} + +#[derive(Clone)] +struct SchemeReplayClient { + scheme: S, + responses: Arc>>, + stream_calls: Arc, +} + +impl SchemeReplayClient { + fn new(scheme: S, responses: Vec>) -> Self { + Self { + scheme, + responses: Arc::new(responses), + stream_calls: Arc::new(AtomicUsize::new(0)), + } + } + + fn stream_count(&self) -> usize { + self.stream_calls.load(Ordering::SeqCst) + } +} + +#[async_trait] +impl LlmClient for SchemeReplayClient { + async fn stream(&self, _request: Request) -> Result { + let response_index = self.stream_calls.fetch_add(1, Ordering::SeqCst); + let frames = self + .responses + .get(response_index) + .ok_or_else(|| ClientError::Config("scheme replay responses exhausted".to_string()))?; + let mut state = S::State::default(); + let mut output = Vec::new(); + for frame in frames { + match self + .scheme + .parse_sse(frame.event_type, frame.data, &mut state) + { + Ok(events) => output.extend(events.into_iter().map(Ok)), + Err(error) => { + output.push(Err(error)); + break; + } + } + } + Ok(Box::pin(futures::stream::iter(output))) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + fn tool_call_completion_support(&self) -> ToolCallCompletionSupport { + self.scheme.tool_call_completion_support() + } +} + +async fn assert_response_complete_scheme_accepts_clean_eof( + scheme: S, + frames: Vec, +) { + assert_eq!( + scheme.tool_call_completion_support(), + ToolCallCompletionSupport::ResponseComplete + ); + + for mode in [ + ToolCallDispatchMode::AfterResponse, + ToolCallDispatchMode::OnToolCallComplete, + ] { + let client = SchemeReplayClient::new(scheme.clone(), vec![frames.clone()]); + let probe = client.clone(); + let mut engine = Engine::new(client); + engine.set_tool_call_dispatch_mode(mode); + let mut history = History::new(); + + let output = engine.run(&mut history, "hello").await; + + assert!( + matches!(output.result, EngineRunExit::Finished), + "{mode:?} should accept the response-complete scheme's clean EOF, got {:?}", + output.result + ); + assert_eq!( + probe.stream_count(), + 1, + "{mode:?} must not open a continuation request after the adapter's normal EOF" + ); + } +} + +#[tokio::test] +async fn openai_chat_clean_eof_completes_in_default_and_opt_in_fallback_modes() { + assert_response_complete_scheme_accepts_clean_eof( + OpenAIScheme::new(), + vec![ + SseFrame { + event_type: "message", + data: r#"{"id":"chatcmpl-123","object":"chat.completion.chunk","created":1694268190,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":"stop"}]}"#, + }, + SseFrame { + event_type: "message", + data: "[DONE]", + }, + ], + ) + .await; +} + +#[tokio::test] +async fn gemini_clean_eof_completes_in_default_and_opt_in_fallback_modes() { + assert_response_complete_scheme_accepts_clean_eof( + GeminiScheme::new(), + vec![ + SseFrame { + event_type: "message", + data: r#"{"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"},"finishReason":"STOP","index":0}]}"#, + }, + SseFrame { + event_type: "message", + data: "[DONE]", + }, + ], + ) + .await; +} From e9e0368db49b3f485613384fbf8c4f8318e0b52e Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 14:41:36 +0900 Subject: [PATCH 10/11] fix: preserve external command attachment authority --- crates/server-api/src/lib.rs | 88 ++++- crates/tui/src/backend_worker_picker.rs | 55 ++- crates/workdir/src/workspace.rs | 17 +- crates/workspace-server/src/server.rs | 342 +++++++++++++++++- crates/workspace-server/src/store.rs | 287 +++++++++++++-- openapi/server-api.json | 14 +- web/workspace/deno.json | 2 +- .../src/lib/generated/repository-api.ts | 8 +- .../src/lib/generated/runtime-api.ts | 8 +- .../src/lib/generated/workdir-api.ts | 4 +- .../src/lib/generated/worker-launch-api.ts | 8 +- .../src/lib/workspace/api/workdirs.ts | 41 ++- .../src/lib/workspace/api/workers.test.ts | 42 +++ .../src/lib/workspace/api/workers.ts | 11 +- .../workspace/settings/workdir-permissions.ts | 12 + .../workspace/sidebar/worker-launch.test.ts | 6 +- .../sidebar/worker-workdir-meta.test.ts | 6 +- .../w/[workspaceId]/workers/+page.svelte | 6 +- .../w/[workspaceId]/workers/new/+page.svelte | 4 + web/workspace/tests/workdir-api.test.ts | 52 ++- .../tests/workdir-permissions.test.ts | 22 ++ 21 files changed, 955 insertions(+), 80 deletions(-) create mode 100644 web/workspace/src/lib/workspace/settings/workdir-permissions.ts create mode 100644 web/workspace/tests/workdir-permissions.test.ts diff --git a/crates/server-api/src/lib.rs b/crates/server-api/src/lib.rs index 825616bf3..e68a33cf2 100644 --- a/crates/server-api/src/lib.rs +++ b/crates/server-api/src/lib.rs @@ -5083,8 +5083,15 @@ pub struct RuntimeWorkingDirectorySummary { #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] pub enum WorkingDirectorySource { - Repository { repository_key: String }, - ExternalGrant { grant_id: String }, + Repository { + repository_key: String, + }, + ExternalGrant { + grant_id: String, + /// Operator-approved ceiling for this External Workdir grant. Worker + /// attachment authority is projected separately after attenuation. + grant_permissions: ExternalWorkdirPermissions, + }, } /// Public, provider-neutral Workdir inventory projection. @@ -5212,6 +5219,16 @@ impl From for RuntimeWorking } } +impl From for ExternalWorkdirPermissions { + fn from(value: workdir::workspace::WorkdirPermissionSummary) -> Self { + Self { + read: value.read, + write: value.write, + command: value.command, + } + } +} + impl From for WorkingDirectorySummary { fn from(value: workdir::workspace::WorkingDirectorySummary) -> Self { Self { @@ -5221,9 +5238,13 @@ impl From for WorkingDirectorySumma workdir::workspace::WorkingDirectorySource::Repository { repository_key } => { WorkingDirectorySource::Repository { repository_key } } - workdir::workspace::WorkingDirectorySource::ExternalGrant { grant_id } => { - WorkingDirectorySource::ExternalGrant { grant_id } - } + workdir::workspace::WorkingDirectorySource::ExternalGrant { + grant_id, + grant_permissions, + } => WorkingDirectorySource::ExternalGrant { + grant_id, + grant_permissions: grant_permissions.into(), + }, }, creation_selector: value.creation_selector, creation_ref: value.creation_ref, @@ -7977,6 +7998,10 @@ pub struct WorkspaceWorkerDiscoveryPage { #[serde(deny_unknown_fields)] pub struct WorkerWorkdirAttachmentSummary { pub alias: String, + /// Current effective category authority after intersecting the durable link, + /// source grant, and Ticket target ceilings. This is independent from + /// provider liveness and may be narrower than the grant ceiling. + pub effective_permissions: ExternalWorkdirPermissions, pub working_directory: WorkingDirectorySummary, } @@ -11269,6 +11294,59 @@ mod tests { } } + #[test] + fn external_workdir_summary_separates_grant_and_attachment_effective_permissions() { + let summary = WorkingDirectorySummary { + working_directory_id: "external-1".into(), + display_name: Some("Session data".into()), + source: WorkingDirectorySource::ExternalGrant { + grant_id: "grant-1".into(), + grant_permissions: ExternalWorkdirPermissions { + read: true, + write: false, + command: true, + }, + }, + creation_selector: None, + creation_ref: None, + creation_tree: None, + current_selector: None, + current_ref: None, + current_tree: None, + observed_at_epoch_seconds: None, + materializer_kind: WorkingDirectoryMaterializerKind::ClientHostedExternal, + cleanup_target: None, + status: WorkingDirectoryStatusKind::Active, + cleanliness: Some("unknown".into()), + occupied_by: None, + }; + let value = serde_json::to_value(WorkerWorkdirAttachmentSummary { + alias: "sessions".into(), + effective_permissions: ExternalWorkdirPermissions { + read: false, + write: false, + command: true, + }, + working_directory: summary, + }) + .unwrap(); + + assert_eq!( + value["working_directory"]["source"]["grant_permissions"], + serde_json::json!({"read": true, "write": false, "command": true}) + ); + assert_eq!( + value["effective_permissions"], + serde_json::json!({"read": false, "write": false, "command": true}) + ); + let mut missing = value; + missing["working_directory"]["source"] + .as_object_mut() + .unwrap() + .remove("grant_permissions"); + assert!(serde_json::from_value::(missing).is_err()); + } + #[test] fn workdir_response_rejects_stale_occupancy_shape() { let stale = serde_json::json!({ diff --git a/crates/tui/src/backend_worker_picker.rs b/crates/tui/src/backend_worker_picker.rs index 4472ea407..8fa950f51 100644 --- a/crates/tui/src/backend_worker_picker.rs +++ b/crates/tui/src/backend_worker_picker.rs @@ -522,14 +522,29 @@ fn working_directory_text(worker: &BackendWorkerSummary) -> String { } let attachment = &worker.workdir_attachments[0]; let directory = &attachment.working_directory; - let source_label = match &directory.source { - BackendWorkingDirectorySource::Repository { repository_key } => repository_key.as_str(), - BackendWorkingDirectorySource::ExternalGrant { .. } => "external", + let (source_label, permission_label) = match &directory.source { + BackendWorkingDirectorySource::Repository { repository_key } => { + (repository_key.as_str(), String::new()) + } + BackendWorkingDirectorySource::ExternalGrant { .. } => { + let permissions = &attachment.effective_permissions; + let mut categories = Vec::new(); + if permissions.read { + categories.push("READ"); + } + if permissions.write { + categories.push("WRITE"); + } + if permissions.command { + categories.push("COMMAND"); + } + ("external", format!("・{}", categories.join("・"))) + } }; let label = directory.display_name.as_deref().unwrap_or(source_label); format!( - "wd:{}:{}・{}", - attachment.alias, label, directory.working_directory_id + "wd:{}:{}{}・{}", + attachment.alias, label, permission_label, directory.working_directory_id ) } @@ -826,6 +841,7 @@ mod tests { worker.workdir_attachments = vec![ serde_json::from_value(serde_json::json!({ "alias": "checkout", + "effective_permissions": {"read": true, "write": true, "command": true}, "working_directory": { "working_directory_id": "001a06a9f0202000000", "display_name": "Checkout", @@ -871,6 +887,7 @@ mod tests { worker.workdir_attachments = vec![ serde_json::from_value(serde_json::json!({ "alias": "checkout", + "effective_permissions": {"read": true, "write": true, "command": true}, "working_directory": { "working_directory_id": "workdir-1", "source": {"kind": "repository", "repository_key": "main"}, @@ -901,6 +918,34 @@ mod tests { ); } + #[test] + fn worker_row_shows_external_attachment_effective_permissions() { + let mut worker = worker("runtime-a", "worker-a", None); + worker.workdir_attachments = vec![ + serde_json::from_value(serde_json::json!({ + "alias": "sessions", + "effective_permissions": {"read": false, "write": false, "command": true}, + "working_directory": { + "working_directory_id": "external-1", + "display_name": "Session data", + "source": { + "kind": "external_grant", + "grant_id": "grant-1", + "grant_permissions": {"read": true, "write": false, "command": true} + }, + "materializer_kind": "client_hosted_external", + "status": "active" + } + })) + .unwrap(), + ]; + + assert_eq!( + working_directory_text(&worker), + "wd:sessions:Session data・COMMAND・external-1" + ); + } + #[test] fn picker_title_uses_backend_worker_wording() { let target = BackendRuntimeListTarget::new( diff --git a/crates/workdir/src/workspace.rs b/crates/workdir/src/workspace.rs index 2e7c57582..8bb21fe06 100644 --- a/crates/workdir/src/workspace.rs +++ b/crates/workdir/src/workspace.rs @@ -100,11 +100,24 @@ pub struct RuntimeWorkingDirectorySummary { pub occupied_by: Option, } +#[derive(Debug, Clone, Copy, Serialize, Deserialize, JsonSchema, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct WorkdirPermissionSummary { + pub read: bool, + pub write: bool, + pub command: bool, +} + #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)] #[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] pub enum WorkingDirectorySource { - Repository { repository_key: String }, - ExternalGrant { grant_id: String }, + Repository { + repository_key: String, + }, + ExternalGrant { + grant_id: String, + grant_permissions: WorkdirPermissionSummary, + }, } #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)] diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 3021be462..2bc9c314c 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -3973,6 +3973,21 @@ fn external_workdir_capabilities(permissions: &str) -> Result server_api::ExternalWorkdirPermissions { + let read = capabilities.supports(workdir::WorkdirSessionCapability::Read) + && capabilities.supports(workdir::WorkdirSessionCapability::Glob) + && capabilities.supports(workdir::WorkdirSessionCapability::Grep); + server_api::ExternalWorkdirPermissions { + read, + write: read + && capabilities.supports(workdir::WorkdirSessionCapability::Write) + && capabilities.supports(workdir::WorkdirSessionCapability::Edit), + command: capabilities.supports(workdir::WorkdirSessionCapability::Command), + } +} + fn workdir_source_capabilities( api: &WorkspaceApi, workdir: &WorkdirRegistryRecord, @@ -4117,6 +4132,19 @@ fn ticket_target_capabilities_for_workdir( Ok(target.capabilities.intersection(source_capabilities)) } +fn effective_worker_workdir_capabilities( + api: &WorkspaceApi, + worker: &RuntimeWorkerRef, + workdir: &WorkdirRegistryRecord, + link_capabilities: workdir::WorkdirSessionCapabilities, +) -> Result { + Ok(link_capabilities + .intersection(workdir_source_capabilities(api, workdir)?) + .intersection(ticket_target_capabilities_for_workdir( + api, worker, workdir, + )?)) +} + #[derive(Clone, Debug)] struct ManualTicketCoderWorkdirBinding { original_links: Vec, @@ -14146,12 +14174,8 @@ async fn open_current_worker_workdir_session_locked( link.workdir_id ), })?; - let effective_capabilities = link - .capabilities - .intersection(workdir_source_capabilities(api, &workdir)?) - .intersection(ticket_target_capabilities_for_workdir( - api, worker, &workdir, - )?); + let effective_capabilities = + effective_worker_workdir_capabilities(api, worker, &workdir, link.capabilities)?; if let WorkdirRegistrySource::ExternalGrant { grant_id } = &workdir.source { close_current_worker_attachment_session_locked(api, worker).await?; let mut grant = api @@ -23752,6 +23776,9 @@ fn project_workspace_worker( api: &WorkspaceApi, summary: InternalWorkerSummary, ) -> ApiResult { + let links = api + .store + .list_worker_workdir_links(&api.config.workspace_id, &summary.worker)?; let workdir_attachments = summary .workdir_attachments .iter() @@ -23761,8 +23788,29 @@ fn project_workspace_worker( &summary.worker.runtime_id, &attachment.working_directory.summary, ); + let link = links + .iter() + .find(|link| { + link.unlinked_at.is_none() + && link.alias == attachment.alias.as_str() + && link.workdir_id == record.workdir_id + }) + .ok_or_else(|| { + Error::RegistryInconsistency(format!( + "Worker {}:{} attachment `{}` has no active durable Workdir link", + summary.worker.runtime_id, summary.worker.worker_id, attachment.alias + )) + })?; Ok(server_api::WorkerWorkdirAttachmentSummary { alias: attachment.alias.to_string(), + effective_permissions: workdir_permission_summary( + effective_worker_workdir_capabilities( + api, + &summary.worker, + &record, + link.capabilities, + )?, + ), working_directory: projected_workdir_summary_from_record(api, &record)?.into(), }) }) @@ -24720,10 +24768,14 @@ fn project_worker_registry_projection( else { continue; }; + let effective_permissions = workdir_permission_summary( + effective_worker_workdir_capabilities(api, &record.worker, workdir, link.capabilities)?, + ); workdir_attachments.push(server_api::WorkerWorkdirAttachmentSummary { alias: workdir::WorkdirAttachmentAlias::new(link.alias) .expect("persisted attachment aliases are validated on write") .to_string(), + effective_permissions, working_directory: projected_workdir_summary_from_record(api, workdir)?.into(), }); } @@ -25106,6 +25158,7 @@ fn runtime_workdir_summary_from_record( fn workdir_summary_from_record( record: &WorkdirRegistryRecord, repository_key: Option<&str>, + external_grant_permissions: Option, ) -> WorkingDirectorySummary { let status = match record.materialization_status.as_str() { "present" => WorkingDirectoryStatusKind::Active, @@ -25134,6 +25187,8 @@ fn workdir_summary_from_record( WorkdirRegistrySource::ExternalGrant { grant_id } => ( WorkingDirectorySource::ExternalGrant { grant_id: grant_id.clone(), + grant_permissions: external_grant_permissions + .expect("External Workdir projection requires grant permissions"), }, MaterializerKind::ClientHostedExternal, None, @@ -25192,16 +25247,39 @@ fn projected_workdir_summary_from_record( api: &WorkspaceApi, record: &WorkdirRegistryRecord, ) -> Result { - let repository_key = match &record.source { - WorkdirRegistrySource::Repository { repository_id, .. } => Some( - api.store - .get_repository(&record.workspace_id, repository_id)? - .ok_or_else(|| Error::UnknownRepository(repository_id.clone()))? - .repository_key, + let (repository_key, external_grant_permissions) = match &record.source { + WorkdirRegistrySource::Repository { repository_id, .. } => ( + Some( + api.store + .get_repository(&record.workspace_id, repository_id)? + .ok_or_else(|| Error::UnknownRepository(repository_id.clone()))? + .repository_key, + ), + None, + ), + WorkdirRegistrySource::ExternalGrant { grant_id } => ( + None, + Some({ + let permissions = api + .store + .get_external_workdir_grant(&record.workspace_id, grant_id)? + .ok_or_else(|| { + Error::Store("External Workdir grant is unavailable".to_string()) + }) + .and_then(|grant| external_workdir_permissions(&grant.permissions))?; + workdir::workspace::WorkdirPermissionSummary { + read: permissions.read, + write: permissions.write, + command: permissions.command, + } + }), ), - WorkdirRegistrySource::ExternalGrant { .. } => None, }; - let mut summary = workdir_summary_from_record(record, repository_key.as_deref()); + let mut summary = workdir_summary_from_record( + record, + repository_key.as_deref(), + external_grant_permissions, + ); apply_workdir_occupancy_projection(api, &mut summary)?; Ok(summary) } @@ -27375,13 +27453,25 @@ mod tests { let mut workdir_ids = Vec::new(); for index in 1..=2 { + let (permissions, capabilities) = if index == 1 { + ( + "command_only", + workdir::WorkdirSessionCapabilities::COMMAND_ONLY, + ) + } else { + ( + "read_command", + workdir::WorkdirSessionCapabilities::READ_ONLY + .union(workdir::WorkdirSessionCapabilities::COMMAND_ONLY), + ) + }; let grant = ExternalWorkdirGrantRecord { grant_id: format!("launch-grant-{index}"), workspace_id: TEST_WORKSPACE_ID.to_string(), workdir_id: format!("launch-external-{index}"), provider_instance_id: format!("launch-provider-{index}"), display_name: format!("External launch {index}"), - permissions: "read_only".to_string(), + permissions: permissions.to_string(), created_by: test_browser_request_actor().account_id, created_at: now.clone(), expires_at: Some(expires_at.to_rfc3339_opts(SecondsFormat::Nanos, true)), @@ -27432,7 +27522,7 @@ mod tests { provider_instance_id: grant.provider_instance_id.clone(), generation: 1, expires_at: Some(expires_at), - capabilities: workdir::WorkdirSessionCapabilities::READ_ONLY, + capabilities, admission: Arc::new(tokio::sync::Semaphore::new(16)), shutdown_confirmed: Arc::new(std::sync::atomic::AtomicBool::new(false)), sender, @@ -27454,6 +27544,33 @@ mod tests { .map(|summary| summary.working_directory_id.as_str()) .collect::>(); assert!(workdir_ids.iter().all(|id| available.contains(id.as_str()))); + for (index, workdir_id) in workdir_ids.iter().enumerate() { + let summary = options + .working_directories + .iter() + .find(|summary| summary.working_directory_id == *workdir_id) + .unwrap(); + let expected_permissions = if index == 0 { + server_api::ExternalWorkdirPermissions { + read: false, + write: false, + command: true, + } + } else { + server_api::ExternalWorkdirPermissions { + read: true, + write: false, + command: true, + } + }; + assert!(matches!( + &summary.source, + server_api::WorkingDirectorySource::ExternalGrant { + grant_permissions, + .. + } if *grant_permissions == expected_permissions + )); + } let Json(worker) = create_workspace_worker( State(api.clone()), @@ -27491,7 +27608,194 @@ mod tests { .collect::>(), HashSet::from(["primary", "reference"]), ); - drop(receivers); + + let worker_ref = RuntimeWorkerRef::new(&worker.runtime_id, &worker.worker_id); + let links = api + .store + .list_worker_workdir_links(TEST_WORKSPACE_ID, &worker_ref) + .unwrap(); + let projected_worker = workers_response(api.clone()) + .unwrap() + .items + .into_iter() + .find(|candidate| candidate.worker_id == worker.worker_id) + .unwrap(); + for (index, receiver) in receivers.iter_mut().enumerate() { + let expected_capabilities = if index == 0 { + workdir::WorkdirSessionCapabilities::COMMAND_ONLY + } else { + workdir::WorkdirSessionCapabilities::READ_ONLY + .union(workdir::WorkdirSessionCapabilities::COMMAND_ONLY) + }; + let link = links + .iter() + .find(|link| link.workdir_id == workdir_ids[index]) + .unwrap(); + assert_eq!(link.capabilities, expected_capabilities); + let projected_attachment = projected_worker + .workdir_attachments + .iter() + .find(|attachment| { + attachment.working_directory.working_directory_id == workdir_ids[index] + }) + .unwrap(); + assert_eq!( + projected_attachment.effective_permissions, + workdir_permission_summary(expected_capabilities) + ); + + let session_lock = current_worker_session_lock(&api, &worker_ref); + let session_guard = session_lock.lock().await; + let session = open_current_worker_workdir_session_locked(&api, &worker_ref, link) + .await + .unwrap(); + drop(session_guard); + assert_eq!(session.capabilities(), expected_capabilities); + if index == 0 { + let denied = session + .read(workdir::ReadRequest { + path: workdir::WorkdirPath::new("forbidden.txt").unwrap(), + offset: 0, + limit: 1, + max_bytes: 1, + }) + .await; + assert!(matches!( + denied, + Err(workdir::WorkdirError::Unsupported( + workdir::WorkdirSessionCapability::Read + )) | Err(workdir::WorkdirError::Denied(_)) + )); + } + + let mut starter = { + let session = session.clone(); + tokio::spawn(async move { + session + .start_command(workdir::CommandRequest { + command: "printf attached".to_string(), + timeout_secs: 30, + output_limit: 1024, + cwd: workdir::WorkdirPath::root(), + spill_dir: None, + tool_call_id: None, + }) + .await + }) + }; + let provider_handle = workdir::CommandHandle(format!("provider-command-{index}")); + let provider_operation = tokio::select! { + result = &mut starter => panic!("attached command ended before provider routing: {result:?}"), + operation = receiver.recv() => operation, + _ = tokio::time::sleep(std::time::Duration::from_secs(2)) => { + panic!("attached command was not routed to the provider") + } + }; + let Some(ExternalProviderCommand::Operation { + operation, + response, + .. + }) = provider_operation + else { + panic!("expected attached command start operation"); + }; + assert!(matches!( + operation, + WorkdirSessionOperation::CommandStart(_) + )); + response + .send(Ok(WorkdirSessionOperationResult::CommandStart( + provider_handle.clone(), + ))) + .unwrap(); + let external_handle = tokio::time::timeout(std::time::Duration::from_secs(2), starter) + .await + .expect("attached command start response timed out") + .unwrap() + .unwrap(); + + let output_handle = external_handle.clone(); + let mut canceler = { + let session = session.clone(); + tokio::spawn(async move { session.cancel_command(external_handle).await }) + }; + let provider_cancellation = tokio::select! { + result = &mut canceler => panic!("attached cancellation ended before provider routing: {result:?}"), + operation = receiver.recv() => operation, + _ = tokio::time::sleep(std::time::Duration::from_secs(2)) => { + panic!("attached cancellation was not routed to the provider") + } + }; + let Some(ExternalProviderCommand::Operation { + operation, + response, + .. + }) = provider_cancellation + else { + panic!("expected attached command cancellation"); + }; + assert!(matches!( + operation, + WorkdirSessionOperation::CommandCancel(ref handle) if handle == &provider_handle + )); + response + .send(Ok(WorkdirSessionOperationResult::CommandCancel)) + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), canceler) + .await + .expect("attached command cancellation response timed out") + .unwrap() + .unwrap(); + + let output = { + let session = session.clone(); + tokio::spawn(async move { + session + .command_output(workdir::CommandOutputRequest { + handle: output_handle, + cursor: 0, + limit: 1024, + wait: true, + }) + .await + }) + }; + let Some(ExternalProviderCommand::Operation { + operation, + response, + .. + }) = tokio::time::timeout(std::time::Duration::from_secs(2), receiver.recv()) + .await + .expect("attached terminal output was not routed to the provider") + else { + panic!("expected attached terminal command output"); + }; + assert!(matches!( + operation, + WorkdirSessionOperation::CommandOutput(workdir::CommandOutputRequest { + wait: true, + .. + }) + )); + 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, + }, + ))) + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(2), output) + .await + .expect("attached terminal output response timed out") + .unwrap() + .unwrap(); + } } #[tokio::test] @@ -29751,7 +30055,7 @@ mod tests { assert_eq!(attachment.alias, "attachment"); assert!(matches!( &attachment.working_directory.source, - server_api::WorkingDirectorySource::ExternalGrant { grant_id } + server_api::WorkingDirectorySource::ExternalGrant { grant_id, .. } if grant_id == "external-grant" )); assert_eq!( @@ -32031,7 +32335,7 @@ mod tests { updated_at: "2".to_string(), }; - let projected = workdir_summary_from_record(&workdir, Some("main")); + let projected = workdir_summary_from_record(&workdir, Some("main"), None); assert_eq!(projected.status, WorkingDirectoryStatusKind::Active); let serialized = serde_json::to_string(&projected).unwrap(); diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 9c902f5b3..1d8769f11 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -21,7 +21,7 @@ use crate::workspace_deletion::WorkspaceDeletionStore; use crate::{Error, Result}; const OLDEST_SCHEMA_VERSION: i64 = 50; -const LATEST_SCHEMA_VERSION: i64 = 68; +const LATEST_SCHEMA_VERSION: i64 = 69; const SCHEMA_BASELINE_NAME: &str = "workspace schema baseline"; const WORKSPACE_RUNTIME_BINDINGS_MIGRATION_NAME: &str = "workspace runtime bindings"; const RUNTIME_BINDING_AUDIT_MIGRATION_NAME: &str = "workspace Runtime binding revision and audit"; @@ -52,6 +52,8 @@ const EXTERNAL_WORKDIR_ACCESS_AND_OPTIONAL_EXPIRY_MIGRATION_NAME: &str = "External Workdir read-write access and optional expiry"; const EXTERNAL_WORKDIR_COMMAND_PERMISSION_MIGRATION_NAME: &str = "External Workdir explicit command permission"; +const WORKDIR_COMMAND_ATTACHMENT_CAPABILITIES_MIGRATION_NAME: &str = + "External Workdir command attachment capabilities"; const TICKET_SCHEMA_VERSION_WITH_TARGETS: i64 = 7; const TICKET_SCHEMA_BASELINE_NAME: &str = "ticket schema baseline"; const WORKER_REGISTRY_PROJECTION_SCHEMA: &str = r#" @@ -189,6 +191,11 @@ const MIGRATIONS: &[Migration] = &[ name: EXTERNAL_WORKDIR_COMMAND_PERMISSION_MIGRATION_NAME, apply: migrate_external_workdir_command_permission_v67_to_v68, }, + Migration { + version: 69, + name: WORKDIR_COMMAND_ATTACHMENT_CAPABILITIES_MIGRATION_NAME, + apply: migrate_workdir_command_attachment_capabilities_v68_to_v69, + }, ]; #[derive(Clone, Copy)] @@ -9138,37 +9145,52 @@ fn read_device_login_flow_record( }) } +const WORKDIR_LINK_CAPABILITY_ENCODINGS: &[(workdir::WorkdirSessionCapabilities, &str)] = &[ + (workdir::WorkdirSessionCapabilities::ALL, "all"), + (workdir::WorkdirSessionCapabilities::READ_ONLY, "read_only"), + ( + workdir::WorkdirSessionCapabilities::READ_WRITE, + "read_write", + ), + ( + workdir::WorkdirSessionCapabilities::COMMAND_ONLY, + "command_only", + ), + ( + workdir::WorkdirSessionCapabilities::READ_ONLY + .union(workdir::WorkdirSessionCapabilities::COMMAND_ONLY), + "read_command", + ), +]; + fn encode_workdir_link_capabilities( capabilities: workdir::WorkdirSessionCapabilities, ) -> Result<&'static str> { - if capabilities == workdir::WorkdirSessionCapabilities::ALL { - Ok("all") - } else if capabilities == workdir::WorkdirSessionCapabilities::READ_ONLY { - Ok("read_only") - } else if capabilities == workdir::WorkdirSessionCapabilities::READ_WRITE { - Ok("read_write") - } else { - Err(Error::InvalidInput( - "Workdir attachment capabilities must be exactly all, read_only, or read_write" - .to_string(), - )) - } + WORKDIR_LINK_CAPABILITY_ENCODINGS + .iter() + .find_map(|(candidate, encoded)| (*candidate == capabilities).then_some(*encoded)) + .ok_or_else(|| { + Error::InvalidInput( + "Workdir attachment capabilities must be exactly all, read_only, read_write, command_only, or read_command" + .to_string(), + ) + }) } fn decode_workdir_link_capabilities( value: &str, column: usize, ) -> rusqlite::Result { - match value { - "all" => Ok(workdir::WorkdirSessionCapabilities::ALL), - "read_only" => Ok(workdir::WorkdirSessionCapabilities::READ_ONLY), - "read_write" => Ok(workdir::WorkdirSessionCapabilities::READ_WRITE), - _ => Err(rusqlite::Error::FromSqlConversionFailure( - column, - rusqlite::types::Type::Text, - format!("invalid Worker Workdir link capabilities {value:?}").into(), - )), - } + WORKDIR_LINK_CAPABILITY_ENCODINGS + .iter() + .find_map(|(capabilities, encoded)| (*encoded == value).then_some(*capabilities)) + .ok_or_else(|| { + rusqlite::Error::FromSqlConversionFailure( + column, + rusqlite::types::Type::Text, + format!("invalid Worker Workdir link capabilities {value:?}").into(), + ) + }) } fn read_worker_workdir_link_record( @@ -12189,6 +12211,103 @@ fn migrate_external_workdir_command_permission_v67_to_v68(conn: &Connection) -> } } +fn migrate_workdir_command_attachment_capabilities_v68_to_v69(conn: &Connection) -> Result<()> { + let current = current_schema_version(conn)?; + if current != 68 { + return Err(Error::Store(format!( + "expected schema version 68 before {WORKDIR_COMMAND_ATTACHMENT_CAPABILITIES_MIGRATION_NAME} migration, found {current}" + ))); + } + let foreign_keys_enabled = + conn.query_row("PRAGMA foreign_keys", [], |row| row.get::<_, i64>(0))? != 0; + if foreign_keys_enabled { + conn.execute_batch("PRAGMA foreign_keys = OFF;")?; + } + let result = (|| { + let tx = rusqlite::Transaction::new_unchecked(conn, TransactionBehavior::Exclusive)?; + tx.execute_batch( + r#" + CREATE TABLE worker_workdir_links_v69 ( + workspace_id TEXT NOT NULL, + runtime_id TEXT NOT NULL, + worker_id TEXT NOT NULL, + workdir_id TEXT NOT NULL, + alias TEXT NOT NULL, + capabilities TEXT NOT NULL CHECK (capabilities IN ( + 'all', 'read_only', 'read_write', 'command_only', 'read_command' + )), + linked_at TEXT NOT NULL, + unlinked_at TEXT, + PRIMARY KEY (workspace_id, worker_id, workdir_id, alias), + FOREIGN KEY (workspace_id, worker_id) + REFERENCES worker_registry(workspace_id, worker_id) ON DELETE CASCADE, + FOREIGN KEY (workspace_id, workdir_id) + REFERENCES workdir_registry(workspace_id, workdir_id) ON DELETE CASCADE + ); + INSERT INTO worker_workdir_links_v69 ( + workspace_id, runtime_id, worker_id, workdir_id, alias, capabilities, + linked_at, unlinked_at + ) SELECT workspace_id, runtime_id, worker_id, workdir_id, alias, capabilities, + linked_at, unlinked_at + FROM worker_workdir_links; + DROP TABLE worker_workdir_links; + ALTER TABLE worker_workdir_links_v69 RENAME TO worker_workdir_links; + CREATE UNIQUE INDEX worker_workdir_links_active_workdir_unique + ON worker_workdir_links(workspace_id, workdir_id) WHERE unlinked_at IS NULL; + CREATE UNIQUE INDEX worker_workdir_links_active_alias_unique + ON worker_workdir_links(workspace_id, worker_id, alias) WHERE unlinked_at IS NULL; + CREATE INDEX worker_workdir_links_workdir + ON worker_workdir_links(workspace_id, workdir_id); + CREATE TRIGGER workdir_attachment_insert_blocked_by_runtime_removal + BEFORE INSERT ON worker_workdir_links FOR EACH ROW + WHEN EXISTS ( + SELECT 1 FROM runtime_removal_operations operation + WHERE operation.runtime_id = NEW.runtime_id + AND operation.state IN ('pending', 'cleanup_pending') + ) + BEGIN SELECT RAISE(ABORT, 'runtime_removal_in_progress'); END; + CREATE TRIGGER workdir_attachment_update_blocked_by_runtime_removal + BEFORE UPDATE ON worker_workdir_links FOR EACH ROW + WHEN EXISTS ( + SELECT 1 FROM runtime_removal_operations operation + WHERE operation.runtime_id = NEW.runtime_id + AND operation.state IN ('pending', 'cleanup_pending') + ) + BEGIN SELECT RAISE(ABORT, 'runtime_removal_in_progress'); END; + "#, + )?; + let violations = + tx.query_row("SELECT COUNT(*) FROM pragma_foreign_key_check", [], |row| { + row.get::<_, i64>(0) + })?; + if violations != 0 { + return Err(Error::Store(format!( + "schema-69 migration left {violations} foreign-key violation(s)" + ))); + } + tx.execute( + "INSERT INTO __yoi_schema_migrations (version, name) VALUES (?1, ?2)", + params![ + 69_i64, + WORKDIR_COMMAND_ATTACHMENT_CAPABILITIES_MIGRATION_NAME + ], + )?; + tx.commit()?; + Ok(()) + })(); + let restore = if foreign_keys_enabled { + conn.execute_batch("PRAGMA foreign_keys = ON;") + .map_err(Error::from) + } else { + Ok(()) + }; + match (result, restore) { + (Err(error), _) => Err(error), + (Ok(()), Err(error)) => Err(error), + (Ok(()), Ok(())) => Ok(()), + } +} + fn create_latest_workspace_schema(conn: &Connection) -> Result<()> { conn.execute_batch("DROP TABLE IF EXISTS typed_ticket_targets;")?; conn.execute_batch(include_str!("latest_schema.sql"))?; @@ -12251,7 +12370,9 @@ fn create_latest_workspace_schema(conn: &Connection) -> Result<()> { ON typed_ticket_targets(workspace_id, repository_key, ticket_id); ALTER TABLE worker_workdir_links ADD COLUMN capabilities TEXT NOT NULL - CHECK (capabilities IN ('all', 'read_only', 'read_write')); + CHECK (capabilities IN ( + 'all', 'read_only', 'read_write', 'command_only', 'read_command' + )); "#, )?; ticket::verify_sqlite_ticket_schema(conn).map_err(|error| { @@ -12947,6 +13068,29 @@ mod tests { Ok(()) } + #[test] + fn worker_workdir_link_capability_codec_round_trips_only_canonical_sets() { + for (capabilities, encoded) in WORKDIR_LINK_CAPABILITY_ENCODINGS { + assert_eq!( + encode_workdir_link_capabilities(*capabilities).unwrap(), + *encoded + ); + assert_eq!( + decode_workdir_link_capabilities(encoded, 0).unwrap(), + *capabilities + ); + } + assert!( + encode_workdir_link_capabilities(workdir::WorkdirSessionCapabilities::EMPTY).is_err() + ); + assert!( + encode_workdir_link_capabilities(workdir::WorkdirSessionCapabilities::WRITE_ONLY) + .is_err() + ); + assert!(decode_workdir_link_capabilities("write_only", 0).is_err()); + assert!(decode_workdir_link_capabilities("future_authority", 0).is_err()); + } + #[tokio::test] async fn external_workdir_grant_persists_typed_source_and_fences_generations() { let store = SqliteWorkspaceStore::in_memory().unwrap(); @@ -14328,6 +14472,10 @@ mod tests { version: 68, name: EXTERNAL_WORKDIR_COMMAND_PERMISSION_MIGRATION_NAME.to_string(), }, + WorkspaceSchemaMigrationStep { + version: 69, + name: WORKDIR_COMMAND_ATTACHMENT_CAPABILITIES_MIGRATION_NAME.to_string(), + }, ] ); @@ -14394,6 +14542,10 @@ mod tests { 68, EXTERNAL_WORKDIR_COMMAND_PERMISSION_MIGRATION_NAME.to_string(), ), + ( + 69, + WORKDIR_COMMAND_ATTACHMENT_CAPABILITIES_MIGRATION_NAME.to_string(), + ), ] ); assert!(!table_exists(conn, "trusted_runtime_records")?); @@ -14509,6 +14661,89 @@ mod tests { ); } + #[test] + fn schema_v69_preserves_legacy_attachment_capabilities_and_allows_command_subsets() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("server.db"); + prepare_schema_v50(&path, Some("workspace-a")); + let mut conn = Connection::open(&path).unwrap(); + configure_sqlite(&conn).unwrap(); + for migration in MIGRATIONS + .iter() + .filter(|migration| migration.version <= 68) + { + (migration.apply)(&conn).unwrap(); + } + let tx = conn.transaction().unwrap(); + tx.execute_batch( + r#" + INSERT INTO worker_registry ( + workspace_id, runtime_id, worker_id, display_name, retention_state, + created_at, updated_at + ) VALUES ('workspace-a', 'runtime-a', 'worker-a', 'Worker A', 'normal', '1', '1'); + INSERT INTO external_workdir_grants ( + grant_id, workspace_id, workdir_id, provider_instance_id, display_name, + permissions, created_by, created_at, expires_at, generation, status, updated_at + ) VALUES + ('grant-ro', 'workspace-a', 'external-ro', 'provider-ro', 'Read only', + 'read_only', 'owner', '1', NULL, 1, 'pending', '1'), + ('grant-rw', 'workspace-a', 'external-rw', 'provider-rw', 'Read write', + 'read_write', 'owner', '1', NULL, 1, 'pending', '1'); + INSERT INTO workdir_registry ( + workspace_id, workdir_id, display_name, source_kind, runtime_id, + repository_id, external_grant_id, creation_selector, creation_ref, + materialization_status, cleanliness, created_at, updated_at + ) VALUES + ('workspace-a', 'external-ro', 'Read only', 'external_grant', NULL, + NULL, 'grant-ro', NULL, NULL, 'pending', 'unknown', '1', '1'), + ('workspace-a', 'external-rw', 'Read write', 'external_grant', NULL, + NULL, 'grant-rw', NULL, NULL, 'pending', 'unknown', '1', '1'); + INSERT INTO worker_workdir_links ( + workspace_id, runtime_id, worker_id, workdir_id, alias, capabilities, + linked_at, unlinked_at + ) VALUES + ('workspace-a', 'runtime-a', 'worker-a', 'external-ro', 'readonly', + 'read_only', '1', '2'), + ('workspace-a', 'runtime-a', 'worker-a', 'external-rw', 'readwrite', + 'read_write', '1', '2'); + "#, + ) + .unwrap(); + tx.commit().unwrap(); + + migrate_workdir_command_attachment_capabilities_v68_to_v69(&conn).unwrap(); + assert_eq!(current_schema_version(&conn).unwrap(), 69); + let capabilities = conn + .prepare("SELECT capabilities FROM worker_workdir_links ORDER BY alias") + .unwrap() + .query_map([], |row| row.get::<_, String>(0)) + .unwrap() + .collect::>>() + .unwrap(); + assert_eq!(capabilities, vec!["read_only", "read_write"]); + assert!( + conn.execute( + "UPDATE worker_workdir_links SET capabilities = 'command_only' WHERE alias = 'readonly'", + [], + ) + .is_ok() + ); + assert!( + conn.execute( + "UPDATE worker_workdir_links SET capabilities = 'read_command' WHERE alias = 'readwrite'", + [], + ) + .is_ok() + ); + assert!( + conn.execute( + "UPDATE worker_workdir_links SET capabilities = 'write_only' WHERE alias = 'readonly'", + [], + ) + .is_err() + ); + } + #[test] fn migration_resumes_from_a_valid_partially_applied_chain() { let temp = tempfile::tempdir().unwrap(); @@ -14527,7 +14762,7 @@ mod tests { .map(|migration| migration.version) .collect::>(), vec![ - 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68 + 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69 ] ); SqliteWorkspaceStore::migrate_database(&path).unwrap(); @@ -14536,7 +14771,7 @@ mod tests { current_schema_version(&conn).unwrap(), LATEST_SCHEMA_VERSION ); - assert_eq!(workspace_schema_migration_history(&conn).unwrap().len(), 19); + assert_eq!(workspace_schema_migration_history(&conn).unwrap().len(), 20); assert!(column_exists(&conn, "worker_workdir_links", "alias").unwrap()); assert!(column_exists(&conn, "worker_workdir_links", "capabilities").unwrap()); assert!(!column_exists(&conn, "worker_workdir_links", "role").unwrap()); diff --git a/openapi/server-api.json b/openapi/server-api.json index f04307fe8..85a0e6c25 100644 --- a/openapi/server-api.json +++ b/openapi/server-api.json @@ -11940,12 +11940,17 @@ "alias": { "type": "string" }, + "effective_permissions": { + "$ref": "#/components/schemas/ExternalWorkdirPermissions", + "description": "Current effective category authority after intersecting the durable link,\nsource grant, and Ticket target ceilings. This is independent from\nprovider liveness and may be narrower than the grant ceiling." + }, "working_directory": { "$ref": "#/components/schemas/WorkingDirectorySummary" } }, "required": [ "alias", + "effective_permissions", "working_directory" ], "type": "object" @@ -12235,6 +12240,10 @@ "grant_id": { "type": "string" }, + "grant_permissions": { + "$ref": "#/components/schemas/ExternalWorkdirPermissions", + "description": "Operator-approved ceiling for this External Workdir grant. Worker\nattachment authority is projected separately after attenuation." + }, "kind": { "const": "external_grant", "type": "string" @@ -12242,7 +12251,8 @@ }, "required": [ "kind", - "grant_id" + "grant_id", + "grant_permissions" ], "type": "object" } @@ -14010,7 +14020,7 @@ "info": { "title": "Yoi Server API", "version": "0.1.0", - "x-yoi-source-digest": "sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be" + "x-yoi-source-digest": "sha256:3a298734fd19f3a224e4f1503385fb8a823b612d328224f45f677ebaca28b96c" }, "openapi": "3.1.0", "paths": { diff --git a/web/workspace/deno.json b/web/workspace/deno.json index 05284f684..8836db08e 100644 --- a/web/workspace/deno.json +++ b/web/workspace/deno.json @@ -6,7 +6,7 @@ "dev": "deno run -A npm:vite@7.2.7 dev", "dev:backend": "cd ../.. && cargo run -p yoi-workspace-server --bin yoi-server -- serve --listen 127.0.0.1:8787", "check": "deno run -A npm:@sveltejs/kit@2.49.4 sync && deno run -A npm:svelte-check@4.3.4 --tsconfig ./tsconfig.json", - "test": "deno test --allow-read=src,test,tests,../../docs --allow-env=LOG,VSCODE_TEXTMATE_DEBUG,NODE_ENV tests/workspace-model.test.ts tests/repository-api.test.ts tests/workspace-catalog.test.ts tests/profile-api.test.ts tests/skill-api.test.ts tests/tooltip-contract.test.ts tests/bevel-contract.test.ts tests/bevel-line-contract.test.ts tests/typography-contract.test.ts src/lib/workspace/auth/model.test.ts src/lib/workspace/protocol-validator.test.ts tests/auth-api.test.ts src/lib/workspace/api/http.test.ts src/lib/workspace/api/workers.test.ts src/lib/workspace/api/ticket-browser.test.ts src/lib/workspace/header/breadcrumb-model.test.ts src/lib/workspace/console/chat-submit.test.ts src/lib/workspace/console/composer-delivery.test.ts test/composer-history.test.ts tests/composer-paste.test.ts src/lib/workspace/console/composer-command.test.ts src/lib/workspace/console/composer-draft.test.ts src/lib/workspace/console/composer-completion.test.ts test/console/ansi.test.ts src/lib/workspace/console/model.test.ts src/lib/workspace/console/run-status.test.ts src/lib/workspace/companion/api.test.ts tests/workdir-api.test.ts src/lib/workspace/console/tasks.test.ts test/ticket-detail-route-reuse.test.ts test/spa-shell.test.ts test/objective-pages.test.ts test/repositories/ui.test.ts src/lib/workspace/console/worker-console.ui.test.ts src/lib/workspace/settings/model.test.ts tests/runtime-connection.test.ts tests/runtime-management.test.ts tests/runtime-management-source.test.ts src/lib/workspace/sidebar/override-stack.test.ts src/lib/workspace/sidebar/route-ownership.test.ts src/lib/workspace/sidebar/workers.test.ts src/lib/workspace/sidebar/workspace-switcher.test.ts src/lib/workspace/sidebar/worker-subscription.test.ts src/lib/workspace/sidebar/worker-workdir-meta.test.ts src/lib/workspace/sidebar/worker-launch.test.ts test/sidebar/worker-actions.test.ts src/lib/workspace/tickets/merge-request-resources.test.ts src/lib/workspace/tickets/ticket-panel.test.ts test/merge-request-status.test.ts test/config-source/decodal-grammar.test.ts test/config-source/editor-state.test.ts test/config-source/fixed-schema-wrapper.test.ts test/config-source/toolchain.test.ts test/config-source/wasm-parity.test.ts test/config-source/api.test.ts tests/memory-api.test.ts test/repository-access/api.test.ts test/repository-access/loader.test.ts test/repository-access/model.test.ts test/repository-access/structure.test.ts test/repository-ssh-connection-ui.test.ts", + "test": "deno test --allow-read=src,test,tests,../../docs --allow-env=LOG,VSCODE_TEXTMATE_DEBUG,NODE_ENV tests/workspace-model.test.ts tests/repository-api.test.ts tests/workspace-catalog.test.ts tests/profile-api.test.ts tests/skill-api.test.ts tests/tooltip-contract.test.ts tests/bevel-contract.test.ts tests/bevel-line-contract.test.ts tests/typography-contract.test.ts src/lib/workspace/auth/model.test.ts src/lib/workspace/protocol-validator.test.ts tests/auth-api.test.ts src/lib/workspace/api/http.test.ts src/lib/workspace/api/workers.test.ts src/lib/workspace/api/ticket-browser.test.ts src/lib/workspace/header/breadcrumb-model.test.ts src/lib/workspace/console/chat-submit.test.ts src/lib/workspace/console/composer-delivery.test.ts test/composer-history.test.ts tests/composer-paste.test.ts src/lib/workspace/console/composer-command.test.ts src/lib/workspace/console/composer-draft.test.ts src/lib/workspace/console/composer-completion.test.ts test/console/ansi.test.ts src/lib/workspace/console/model.test.ts src/lib/workspace/console/run-status.test.ts src/lib/workspace/companion/api.test.ts tests/workdir-api.test.ts tests/workdir-permissions.test.ts src/lib/workspace/console/tasks.test.ts test/ticket-detail-route-reuse.test.ts test/spa-shell.test.ts test/objective-pages.test.ts test/repositories/ui.test.ts src/lib/workspace/console/worker-console.ui.test.ts src/lib/workspace/settings/model.test.ts tests/runtime-connection.test.ts tests/runtime-management.test.ts tests/runtime-management-source.test.ts src/lib/workspace/sidebar/override-stack.test.ts src/lib/workspace/sidebar/route-ownership.test.ts src/lib/workspace/sidebar/workers.test.ts src/lib/workspace/sidebar/workspace-switcher.test.ts src/lib/workspace/sidebar/worker-subscription.test.ts src/lib/workspace/sidebar/worker-workdir-meta.test.ts src/lib/workspace/sidebar/worker-launch.test.ts test/sidebar/worker-actions.test.ts src/lib/workspace/tickets/merge-request-resources.test.ts src/lib/workspace/tickets/ticket-panel.test.ts test/merge-request-status.test.ts test/config-source/decodal-grammar.test.ts test/config-source/editor-state.test.ts test/config-source/fixed-schema-wrapper.test.ts test/config-source/toolchain.test.ts test/config-source/wasm-parity.test.ts test/config-source/api.test.ts tests/memory-api.test.ts test/repository-access/api.test.ts test/repository-access/loader.test.ts test/repository-access/model.test.ts test/repository-access/structure.test.ts test/repository-ssh-connection-ui.test.ts", "test:component": "deno run -A npm:vitest@3.2.4 run --config vitest.config.ts", "build": "deno run -A npm:vite@7.2.7 build", "preview": "deno run -A npm:vite@7.2.7 preview", diff --git a/web/workspace/src/lib/generated/repository-api.ts b/web/workspace/src/lib/generated/repository-api.ts index cb5c39c72..04e182360 100644 --- a/web/workspace/src/lib/generated/repository-api.ts +++ b/web/workspace/src/lib/generated/repository-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-repository-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) +// Input contract: 0.1.0 (sha256:3a298734fd19f3a224e4f1503385fb8a823b612d328224f45f677ebaca28b96c) // Options: repository-operation-schema-closure,openapi-3.1,format:uint64=safe-number,strict-safe-integers,sorted-output // Output: web/workspace/src/lib/generated/repository-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_repository_openapi_types -- --check @@ -15,6 +15,8 @@ export type Diagnostic = { code: string; message: string; severity: DiagnosticSe export type DiagnosticSeverity = "info" | "warning" | "error"; +export type ExternalWorkdirPermissions = { command: boolean; read: boolean; write: boolean; }; + export type GitCommitSummary = { author_date: string; author_email: string; author_name: string; hash: string; parents: Array; refs: Array; short_hash: string; summary: string; }; export type GitRemoteSummary = { fetch_url: string; name: string; }; @@ -73,7 +75,7 @@ export type WorkerStateSnapshot = { last_command_id: number; state: WorkerState; export type WorkerSummary = { availability: SubscriptionWorkerAvailability; capabilities: WorkerCapabilitySummary; diagnostics?: Array; display_name?: string; host_id: string; implementation: WorkerImplementationSummary; label: string; last_seen_at?: string | null; pinned?: boolean; profile?: string | null; resource_key: string; retention_state?: string; runtime_id: string; singleton_key?: string | null; state: string; tags?: Array; workdir_attachments?: Array; worker_id: string; worker_state?: WorkerStateSnapshot | null; workspace: WorkerWorkspaceSummary; }; -export type WorkerWorkdirAttachmentSummary = { alias: string; working_directory: WorkingDirectorySummary; }; +export type WorkerWorkdirAttachmentSummary = { alias: string; effective_permissions: ExternalWorkdirPermissions; working_directory: WorkingDirectorySummary; }; export type WorkerWorkspaceSummary = { identity: string; visibility: string; workspace_id?: string | null; }; @@ -83,7 +85,7 @@ export type WorkingDirectoryMaterializerKind = "runtime_git_clone" | "client_hos export type WorkingDirectoryOccupancy = { display_name: string; linked_at: string; runtime_id: string; worker_id: string; }; -export type WorkingDirectorySource = { kind: "repository"; repository_key: string; } | { grant_id: string; kind: "external_grant"; }; +export type WorkingDirectorySource = { kind: "repository"; repository_key: string; } | { grant_id: string; grant_permissions: ExternalWorkdirPermissions; kind: "external_grant"; }; export type WorkingDirectoryStatusKind = "active" | "cleanup_pending" | "corrupted" | "not_found" | "unknown"; diff --git a/web/workspace/src/lib/generated/runtime-api.ts b/web/workspace/src/lib/generated/runtime-api.ts index 65d6017f2..83d4c91a0 100644 --- a/web/workspace/src/lib/generated/runtime-api.ts +++ b/web/workspace/src/lib/generated/runtime-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-runtime-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) +// Input contract: 0.1.0 (sha256:3a298734fd19f3a224e4f1503385fb8a823b612d328224f45f677ebaca28b96c) // Options: runtime-operation-schema-closure,openapi-3.1,strict-integer-ranges,sorted-output // Output: web/workspace/src/lib/generated/runtime-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_runtime_api_types -- --check @@ -27,6 +27,8 @@ export type DiagnosticSeverity = "info" | "warning" | "error"; export type ExecuteRuntimeCleanupRequest = { confirm_dirty_discard_target_ids?: Array; expected_plan_digest: string; expected_plan_revision: string; workdir_target_ids?: Array; worker_target_ids?: Array; }; +export type ExternalWorkdirPermissions = { command: boolean; read: boolean; write: boolean; }; + export type RemoveRuntimeRequest = { expected_binding_revision: number; operation_id: string; }; export type RepositoryApiError = { diagnostics?: Array; error: string; message: string; }; @@ -163,7 +165,7 @@ export type WorkerStateSnapshot = { last_command_id: number; state: WorkerState; export type WorkerSummary = { availability: SubscriptionWorkerAvailability; capabilities: WorkerCapabilitySummary; diagnostics?: Array; display_name?: string; host_id: string; implementation: WorkerImplementationSummary; label: string; last_seen_at?: string | null; pinned?: boolean; profile?: string | null; resource_key: string; retention_state?: string; runtime_id: string; singleton_key?: string | null; state: string; tags?: Array; workdir_attachments?: Array; worker_id: string; worker_state?: WorkerStateSnapshot | null; workspace: WorkerWorkspaceSummary; }; -export type WorkerWorkdirAttachmentSummary = { alias: string; working_directory: WorkingDirectorySummary; }; +export type WorkerWorkdirAttachmentSummary = { alias: string; effective_permissions: ExternalWorkdirPermissions; working_directory: WorkingDirectorySummary; }; export type WorkerWorkspaceSummary = { identity: string; visibility: string; workspace_id?: string | null; }; @@ -173,7 +175,7 @@ export type WorkingDirectoryMaterializerKind = "runtime_git_clone" | "client_hos export type WorkingDirectoryOccupancy = { display_name: string; linked_at: string; runtime_id: string; worker_id: string; }; -export type WorkingDirectorySource = { kind: "repository"; repository_key: string; } | { grant_id: string; kind: "external_grant"; }; +export type WorkingDirectorySource = { kind: "repository"; repository_key: string; } | { grant_id: string; grant_permissions: ExternalWorkdirPermissions; kind: "external_grant"; }; export type WorkingDirectoryStatusKind = "active" | "cleanup_pending" | "corrupted" | "not_found" | "unknown"; diff --git a/web/workspace/src/lib/generated/workdir-api.ts b/web/workspace/src/lib/generated/workdir-api.ts index ed5c7392e..c784459f9 100644 --- a/web/workspace/src/lib/generated/workdir-api.ts +++ b/web/workspace/src/lib/generated/workdir-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-workdir-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) +// Input contract: 0.1.0 (sha256:3a298734fd19f3a224e4f1503385fb8a823b612d328224f45f677ebaca28b96c) // Options: workdir-operation-schema-closure,openapi-3.1,strict-integer-ranges,sorted-output // Output: web/workspace/src/lib/generated/workdir-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_workdir_api_types -- --check @@ -37,7 +37,7 @@ export type WorkingDirectoryRemovalRequest = { reason: string; }; export type WorkingDirectoryRemovalResponse = { disposition: WorkingDirectoryRemovalDisposition; failure_category?: string | null; retryable: boolean; working_directory_id: string; }; -export type WorkingDirectorySource = { kind: "repository"; repository_key: string; } | { grant_id: string; kind: "external_grant"; }; +export type WorkingDirectorySource = { kind: "repository"; repository_key: string; } | { grant_id: string; grant_permissions: ExternalWorkdirPermissions; kind: "external_grant"; }; export type WorkingDirectoryStatusKind = "active" | "cleanup_pending" | "corrupted" | "not_found" | "unknown"; diff --git a/web/workspace/src/lib/generated/worker-launch-api.ts b/web/workspace/src/lib/generated/worker-launch-api.ts index 2e2c07ed4..ce3cb7aa0 100644 --- a/web/workspace/src/lib/generated/worker-launch-api.ts +++ b/web/workspace/src/lib/generated/worker-launch-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-worker-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) +// Input contract: 0.1.0 (sha256:3a298734fd19f3a224e4f1503385fb8a823b612d328224f45f677ebaca28b96c) // Options: worker-operation-schema-closure,openapi-3.1,strict-integer-ranges,sorted-output // Output: web/workspace/src/lib/generated/worker-launch-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_worker_launch_api_types -- --check @@ -19,6 +19,8 @@ export type Diagnostic = { code: string; message: string; severity: DiagnosticSe export type DiagnosticSeverity = "info" | "warning" | "error"; +export type ExternalWorkdirPermissions = { command: boolean; read: boolean; write: boolean; }; + export type PasteArtifactAvailability = "available" | "unavailable" | "integrity_failed"; export type PasteArtifactMediaType = "text_plain_utf8"; @@ -67,7 +69,7 @@ export type WorkerStateSnapshot = { last_command_id: number; state: WorkerState; export type WorkerSummary = { availability: SubscriptionWorkerAvailability; capabilities: WorkerCapabilitySummary; diagnostics?: Array; display_name?: string; host_id: string; implementation: WorkerImplementationSummary; label: string; last_seen_at?: string | null; pinned?: boolean; profile?: string | null; resource_key: string; retention_state?: string; runtime_id: string; singleton_key?: string | null; state: string; tags?: Array; workdir_attachments?: Array; worker_id: string; worker_state?: WorkerStateSnapshot | null; workspace: WorkerWorkspaceSummary; }; -export type WorkerWorkdirAttachmentSummary = { alias: string; working_directory: WorkingDirectorySummary; }; +export type WorkerWorkdirAttachmentSummary = { alias: string; effective_permissions: ExternalWorkdirPermissions; working_directory: WorkingDirectorySummary; }; export type WorkerWorkspaceSummary = { identity: string; visibility: string; workspace_id?: string | null; }; @@ -79,7 +81,7 @@ export type WorkingDirectoryOccupancy = { display_name: string; linked_at: strin export type WorkingDirectoryRepositoryOption = { default_selector?: string | null; repository_key: string; }; -export type WorkingDirectorySource = { kind: "repository"; repository_key: string; } | { grant_id: string; kind: "external_grant"; }; +export type WorkingDirectorySource = { kind: "repository"; repository_key: string; } | { grant_id: string; grant_permissions: ExternalWorkdirPermissions; kind: "external_grant"; }; export type WorkingDirectoryStatusKind = "active" | "cleanup_pending" | "corrupted" | "not_found" | "unknown"; diff --git a/web/workspace/src/lib/workspace/api/workdirs.ts b/web/workspace/src/lib/workspace/api/workdirs.ts index 870c425fd..f22724f07 100644 --- a/web/workspace/src/lib/workspace/api/workdirs.ts +++ b/web/workspace/src/lib/workspace/api/workdirs.ts @@ -1,5 +1,6 @@ import type { Diagnostic, + ExternalWorkdirPermissions, WorkingDirectoryCleanupTarget, WorkingDirectoryCreateRequest, WorkingDirectoryCreateResponse, @@ -176,13 +177,15 @@ export function parseWorkingDirectorySummary( function parseWorkingDirectorySource(value: unknown): WorkingDirectorySource { const source = exactRecord( value, - new Set(["kind", "repository_key", "grant_id"]), + new Set(["kind", "repository_key", "grant_id", "grant_permissions"]), "Workdir source", ); const kind = stringField(source, "kind"); if (kind === "repository") { - if (source.grant_id !== undefined) { - throw new Error("Repository Workdir source must not contain grant_id"); + if (source.grant_id !== undefined || source.grant_permissions !== undefined) { + throw new Error( + "Repository Workdir source must not contain External grant authority", + ); } return { kind, @@ -198,11 +201,31 @@ function parseWorkingDirectorySource(value: unknown): WorkingDirectorySource { return { kind, grant_id: stringField(source, "grant_id"), + grant_permissions: parseExternalWorkdirPermissions( + source.grant_permissions, + "External Workdir grant permissions", + ), }; } throw new Error("Workdir source.kind has an unsupported value"); } +export function parseExternalWorkdirPermissions( + value: unknown, + label = "External Workdir permissions", +): ExternalWorkdirPermissions { + const permissions = exactRecord( + value, + new Set(["read", "write", "command"]), + label, + ); + return { + read: booleanField(permissions, "read", label), + write: booleanField(permissions, "write", label), + command: booleanField(permissions, "command", label), + }; +} + function parseCleanupTarget(value: unknown): WorkingDirectoryCleanupTarget { const record = exactRecord( value, @@ -260,6 +283,18 @@ function stringField(record: Record, key: string): string { return value; } +function booleanField( + record: Record, + key: string, + label: string, +): boolean { + const value = record[key]; + if (typeof value !== "boolean") { + throw new Error(`${label}.${key} must be a boolean`); + } + return value; +} + function arrayField(record: Record, key: string): unknown[] { const value = record[key]; if (!Array.isArray(value)) throw new Error(`${key} must be an array`); diff --git a/web/workspace/src/lib/workspace/api/workers.test.ts b/web/workspace/src/lib/workspace/api/workers.test.ts index 2667a0e63..f6b703159 100644 --- a/web/workspace/src/lib/workspace/api/workers.test.ts +++ b/web/workspace/src/lib/workspace/api/workers.test.ts @@ -96,6 +96,48 @@ Deno.test("Worker summary parser enforces observation freshness", () => { assertEquals(observedStopped.state, "stopped"); }); +Deno.test("Worker summary parser preserves attachment-effective External permissions", () => { + const attachment = { + alias: "sessions", + effective_permissions: { read: false, write: false, command: true }, + working_directory: { + working_directory_id: "external-1", + display_name: "Sessions", + source: { + kind: "external_grant", + grant_id: "grant-1", + grant_permissions: { read: true, write: false, command: true }, + }, + materializer_kind: "client_hosted_external", + status: "active", + }, + }; + const parsed = parseWorkerSummary({ + ...worker, + resource_key: "W-1", + availability: "observed", + workdir_attachments: [attachment], + }); + assertEquals( + parsed.workdir_attachments?.[0]?.effective_permissions, + attachment.effective_permissions, + ); + assertThrows( + () => + parseWorkerSummary({ + ...worker, + resource_key: "W-1", + availability: "observed", + workdir_attachments: [{ + alias: attachment.alias, + working_directory: attachment.working_directory, + }], + }), + Error, + "effective_permissions", + ); +}); + Deno.test("Worker launch options parser accepts the generated wire shape", () => { const parsed = parseWorkerLaunchOptionsResponse({ workspace_id: "workspace-a", diff --git a/web/workspace/src/lib/workspace/api/workers.ts b/web/workspace/src/lib/workspace/api/workers.ts index 26dcefe8c..96ac01d23 100644 --- a/web/workspace/src/lib/workspace/api/workers.ts +++ b/web/workspace/src/lib/workspace/api/workers.ts @@ -24,7 +24,10 @@ import type { WorkingDirectoryRepositoryOption, } from "$lib/generated/worker-launch-api"; import type { Segment } from "$lib/generated/protocol"; -import { parseWorkingDirectorySummary } from "$lib/workspace/api/workdirs"; +import { + parseExternalWorkdirPermissions, + parseWorkingDirectorySummary, +} from "$lib/workspace/api/workdirs"; const DIAGNOSTIC_SEVERITIES = new Set([ "info", @@ -316,9 +319,13 @@ function workerWorkdirAttachment( label: string, ): WorkerWorkdirAttachmentSummary { const item = record(value, label); - exact(item, ["alias", "working_directory"], label); + exact(item, ["alias", "effective_permissions", "working_directory"], label); return { alias: string(item.alias, `${label}.alias`), + effective_permissions: parseExternalWorkdirPermissions( + item.effective_permissions, + `${label}.effective_permissions`, + ), working_directory: parseWorkingDirectorySummary(item.working_directory), }; } diff --git a/web/workspace/src/lib/workspace/settings/workdir-permissions.ts b/web/workspace/src/lib/workspace/settings/workdir-permissions.ts new file mode 100644 index 000000000..7942e324d --- /dev/null +++ b/web/workspace/src/lib/workspace/settings/workdir-permissions.ts @@ -0,0 +1,12 @@ +import type { ExternalWorkdirPermissions } from "$lib/generated/workdir-api"; + +/** Render independent Workdir authority categories in a stable, non-conflating order. */ +export function formatWorkdirPermissions( + permissions: ExternalWorkdirPermissions, +): string { + const categories: string[] = []; + if (permissions.read) categories.push("READ"); + if (permissions.write) categories.push("WRITE"); + if (permissions.command) categories.push("COMMAND"); + return categories.length > 0 ? categories.join(" · ") : "NONE"; +} diff --git a/web/workspace/src/lib/workspace/sidebar/worker-launch.test.ts b/web/workspace/src/lib/workspace/sidebar/worker-launch.test.ts index 9fea0695a..6cd8d19fd 100644 --- a/web/workspace/src/lib/workspace/sidebar/worker-launch.test.ts +++ b/web/workspace/src/lib/workspace/sidebar/worker-launch.test.ts @@ -32,7 +32,11 @@ Deno.test("defaultWorkerLaunchForm preserves an available read-only External Wor const external = { working_directory_id: "external-1", display_name: "Session analysis", - source: { kind: "external_grant" as const, grant_id: "grant-1" }, + source: { + kind: "external_grant" as const, + grant_id: "grant-1", + grant_permissions: { read: true, write: false, command: false }, + }, materializer_kind: "client_hosted_external" as const, status: "active" as const, cleanliness: "unknown", diff --git a/web/workspace/src/lib/workspace/sidebar/worker-workdir-meta.test.ts b/web/workspace/src/lib/workspace/sidebar/worker-workdir-meta.test.ts index 31e320ce2..cd5988819 100644 --- a/web/workspace/src/lib/workspace/sidebar/worker-workdir-meta.test.ts +++ b/web/workspace/src/lib/workspace/sidebar/worker-workdir-meta.test.ts @@ -111,7 +111,11 @@ Deno.test("Worker sidebar Workdir label uses an External Workdir display name wi const meta = sidebarWorkdirMeta([ attachment("sessions", { display_name: "Session analysis", - source: { kind: "external_grant", grant_id: "grant-1" }, + source: { + kind: "external_grant", + grant_id: "grant-1", + grant_permissions: { read: true, write: false, command: false }, + }, materializer_kind: "client_hosted_external", cleanliness: "unknown", }), diff --git a/web/workspace/src/routes/w/[workspaceId]/workers/+page.svelte b/web/workspace/src/routes/w/[workspaceId]/workers/+page.svelte index 0d2c61796..0c84506a3 100644 --- a/web/workspace/src/routes/w/[workspaceId]/workers/+page.svelte +++ b/web/workspace/src/routes/w/[workspaceId]/workers/+page.svelte @@ -7,6 +7,7 @@ parseWorkerRetentionResponse, } from '$lib/workspace/api/runtime-workers'; import { workerHref } from '$lib/workspace/resource-links'; + import { formatWorkdirPermissions } from '$lib/workspace/settings/workdir-permissions'; import { formatCurrentWorkdirRevision } from '$lib/workspace/settings/workdir-revision'; import { canOpenWorkerConsole } from '$lib/workspace/sidebar/workers'; import { liveWorkerState } from '$lib/workspace/sidebar/worker-state'; @@ -149,7 +150,7 @@ function workerDirectory(worker: Worker): string { const attachments = worker.workdir_attachments ?? []; if (attachments.length === 0) return '—'; - return attachments.map(({ alias, working_directory: directory }) => { + return attachments.map(({ alias, effective_permissions, working_directory: directory }) => { const repositoryKey = directory.source.kind === 'repository' ? directory.source.repository_key : null; @@ -157,6 +158,9 @@ ? data.repositories?.items.find((repository) => repository.repository_key === repositoryKey)?.provider : null; const label = directory.display_name ?? repositoryKey ?? 'External Workdir'; + if (directory.source.kind === 'external_grant') { + return `${alias}: ${label} · ${formatWorkdirPermissions(effective_permissions)}`; + } return `${alias}: ${label} · ${formatCurrentWorkdirRevision(directory, provider)}`; }).join(', '); } diff --git a/web/workspace/src/routes/w/[workspaceId]/workers/new/+page.svelte b/web/workspace/src/routes/w/[workspaceId]/workers/new/+page.svelte index db55d1e10..3c2e82aa1 100644 --- a/web/workspace/src/routes/w/[workspaceId]/workers/new/+page.svelte +++ b/web/workspace/src/routes/w/[workspaceId]/workers/new/+page.svelte @@ -12,6 +12,7 @@ parseWorkerLaunchOptionsResponse, } from '$lib/workspace/api/workers'; import { formatCurrentWorkdirRevision } from '$lib/workspace/settings/workdir-revision'; + import { formatWorkdirPermissions } from '$lib/workspace/settings/workdir-permissions'; import { buildCreateWorkspaceWorkerRequest, defaultWorkerLaunchForm, @@ -38,6 +39,9 @@ ? data.repositories?.items.find((repository) => repository.repository_key === repositoryKey)?.provider : null; const label = directory.display_name ?? repositoryKey ?? 'External Workdir'; + if (directory.source.kind === 'external_grant') { + return `${label} · External · ${formatWorkdirPermissions(directory.source.grant_permissions)}`; + } return `${label} · ${formatCurrentWorkdirRevision(directory, provider)}`; } diff --git a/web/workspace/tests/workdir-api.test.ts b/web/workspace/tests/workdir-api.test.ts index f4c125484..abb9d8963 100644 --- a/web/workspace/tests/workdir-api.test.ts +++ b/web/workspace/tests/workdir-api.test.ts @@ -53,7 +53,11 @@ Deno.test("Workdir REST validation accepts External grant sources without Runtim item: { working_directory_id: "external-1", display_name: "Session analysis", - source: { kind: "external_grant", grant_id: "grant-1" }, + source: { + kind: "external_grant", + grant_id: "grant-1", + grant_permissions: { read: true, write: false, command: true }, + }, materializer_kind: "client_hosted_external", status: "active", cleanliness: "unknown", @@ -66,6 +70,52 @@ Deno.test("Workdir REST validation accepts External grant sources without Runtim if (detail.item.source.kind !== "external_grant") { throw new Error("External Workdir source was not preserved"); } + if ( + !detail.item.source.grant_permissions.read || + detail.item.source.grant_permissions.write || + !detail.item.source.grant_permissions.command + ) { + throw new Error("External Workdir grant permissions were not preserved"); + } +}); + +Deno.test("Workdir REST validation rejects missing or misplaced External grant permissions", () => { + for ( + const source of [ + { kind: "external_grant", grant_id: "grant-1" }, + { + kind: "external_grant", + grant_id: "grant-1", + grant_permissions: { read: true, write: false, command: "yes" }, + }, + { + kind: "repository", + repository_key: "main", + grant_permissions: { read: true, write: false, command: false }, + }, + ] + ) { + let rejected = false; + try { + parseWorkingDirectoryDetailResponse({ + workspace_id: "workspace-a", + item: { + working_directory_id: "invalid-1", + source, + materializer_kind: source.kind === "repository" + ? "runtime_git_clone" + : "client_hosted_external", + status: "active", + }, + diagnostics: [], + }); + } catch { + rejected = true; + } + if (!rejected) { + throw new Error(`invalid External permission projection was accepted: ${JSON.stringify(source)}`); + } + } }); Deno.test("Workdir REST validation rejects stale response JSON", () => { diff --git a/web/workspace/tests/workdir-permissions.test.ts b/web/workspace/tests/workdir-permissions.test.ts new file mode 100644 index 000000000..244cb4e47 --- /dev/null +++ b/web/workspace/tests/workdir-permissions.test.ts @@ -0,0 +1,22 @@ +declare const Deno: { + test(name: string, fn: () => void | Promise): void; +}; + +import { formatWorkdirPermissions } from "../src/lib/workspace/settings/workdir-permissions.ts"; + +Deno.test("Workdir permissions render independent categories in stable order", () => { + const cases = [ + [{ read: true, write: false, command: false }, "READ"], + [{ read: true, write: true, command: false }, "READ · WRITE"], + [{ read: false, write: false, command: true }, "COMMAND"], + [{ read: true, write: false, command: true }, "READ · COMMAND"], + [{ read: true, write: true, command: true }, "READ · WRITE · COMMAND"], + ] as const; + + for (const [permissions, expected] of cases) { + const actual = formatWorkdirPermissions(permissions); + if (actual !== expected) { + throw new Error(`expected ${expected}, got ${actual}`); + } + } +}); From 4ffa9e418ed2a61a3c4375ab3334a44047a78b39 Mon Sep 17 00:00:00 2001 From: Hare Date: Tue, 29 Sep 2026 14:49:13 +0900 Subject: [PATCH 11/11] chore: refresh generated API contract hashes --- web/workspace/src/lib/generated/companion-api.ts | 2 +- web/workspace/src/lib/generated/ticket-api.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/web/workspace/src/lib/generated/companion-api.ts b/web/workspace/src/lib/generated/companion-api.ts index 2020f3a12..df244651f 100644 --- a/web/workspace/src/lib/generated/companion-api.ts +++ b/web/workspace/src/lib/generated/companion-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-companion-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) +// Input contract: 0.1.0 (sha256:3a298734fd19f3a224e4f1503385fb8a823b612d328224f45f677ebaca28b96c) // Options: companion-operation-schema-closure,openapi-3.1,strict-integer-ranges,sorted-output // Output: web/workspace/src/lib/generated/companion-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_companion_api_types -- --check diff --git a/web/workspace/src/lib/generated/ticket-api.ts b/web/workspace/src/lib/generated/ticket-api.ts index 2747cd40e..4e8e5ad7d 100644 --- a/web/workspace/src/lib/generated/ticket-api.ts +++ b/web/workspace/src/lib/generated/ticket-api.ts @@ -1,6 +1,6 @@ // @generated by yoi-ticket-browser-openapi-typescript@1. Do not edit. // Input: openapi/server-api.json -// Input contract: 0.1.0 (sha256:e0407f2ae3f1d454d05543627b87eff8eb001b5a90a900ca5af91d7bb492c3be) +// Input contract: 0.1.0 (sha256:3a298734fd19f3a224e4f1503385fb8a823b612d328224f45f677ebaca28b96c) // Options: ticket-objective-merge-request-browser-operation-schema-closure,openapi-3.1,strict-integer-ranges,sorted-output // Output: web/workspace/src/lib/generated/ticket-api.ts // Regenerate/check: cargo run -q -p server-api --example generate_ticket_api_types -- --check