diff --git a/crates/agen/README.md b/crates/agen/README.md index 1a80dacd..13478848 100644 --- a/crates/agen/README.md +++ b/crates/agen/README.md @@ -21,20 +21,21 @@ agen = { version = "0.2.1", features = ["codex"] } ## Quick start -Supply an implementation of [`LlmClient`](https://docs.rs/agen/latest/agen/llm_client/trait.LlmClient.html), then run a turn. The first call consumes the mutable engine and returns a cache-locked engine for later turns. +Supply an implementation of [`LlmClient`](https://docs.rs/agen/latest/agen/llm_client/trait.LlmClient.html), keep conversation history in your application, then run a turn. The first call consumes the mutable engine and returns a cache-locked engine for later turns. ```no_run -use agen::{Engine, EngineError}; +use agen::{Engine, EngineError, History}; use agen::llm_client::LlmClient; async fn conversation(client: C) -> Result<(), EngineError> { + let mut history = History::new(); let output = Engine::new(client) .system_prompt("You are a concise assistant.") - .run("Explain typed state in one sentence.") - .await?; + .run(&mut history, "Explain typed state in one sentence.") + .await; let mut engine = output.engine; - let _result = engine.run("Give a Rust example.").await?; + let _result = engine.run(&mut history, "Give a Rust example.").await; Ok(()) } ``` diff --git a/crates/agen/examples/engine_cancel_demo.rs b/crates/agen/examples/engine_cancel_demo.rs index 6c7cbe13..0ee6f0dc 100644 --- a/crates/agen/examples/engine_cancel_demo.rs +++ b/crates/agen/examples/engine_cancel_demo.rs @@ -4,7 +4,7 @@ use agen::llm_client::scheme::{Scheme, anthropic::AnthropicScheme}; use agen::llm_client::transport::{HttpTransport, ResolvedAuth}; -use agen::{Engine, EngineResult}; +use agen::{Engine, EngineRunExit, StopReason}; use std::time::Duration; #[tokio::main] @@ -29,6 +29,7 @@ async fn main() -> Result<(), Box> { let base_url = scheme.default_base_url().to_string(); let client = HttpTransport::new(scheme, model, base_url, ResolvedAuth::ApiKey(api_key), cap); let engine = Engine::new(client); + let mut history = agen::History::new(); println!("🚀 Starting Engine..."); println!("💡 Will cancel after 2 seconds\n"); @@ -45,16 +46,15 @@ async fn main() -> Result<(), Box> { println!("📡 Sending request to LLM..."); - match engine.run("Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await { - Ok(out) => match out.result { - EngineResult::Finished => println!("✅ Task completed normally"), - EngineResult::Paused => println!("⏸️ Task paused"), - EngineResult::LimitReached => println!("🔒 Turn limit reached"), - EngineResult::Yielded => println!("↩️ Task yielded"), - }, - Err(e) => { - println!("❌ Task error: {}", e); + let output = engine.run(&mut history, "Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await; + match output.result { + EngineRunExit::Finished => println!("✅ Task completed normally"), + EngineRunExit::Paused => println!("⏸️ Task paused"), + EngineRunExit::Yielded => println!("↩️ Task yielded"), + EngineRunExit::Interrupted(StopReason::LimitReached) => { + println!("🔒 Turn limit reached") } + EngineRunExit::Interrupted(reason) => println!("❌ Task interrupted: {reason:?}"), } println!("\n✨ Demo complete!"); diff --git a/crates/agen/examples/engine_cli.rs b/crates/agen/examples/engine_cli.rs index f0ebff9d..9fa2e143 100644 --- a/crates/agen/examples/engine_cli.rs +++ b/crates/agen/examples/engine_cli.rs @@ -39,7 +39,7 @@ use tracing::info; use tracing_subscriber::EnvFilter; use agen::{ - Engine, + Engine, EngineRunExit, StopReason, interceptor::{Interceptor, PostToolAction, ToolResultInfo}, llm_client::{ LlmClient, @@ -451,6 +451,7 @@ async fn main() -> Result<(), Box> { // Create Engine let mut engine = Engine::new(client); + let mut history = agen::History::new(); let tool_call_names = Arc::new(Mutex::new(HashMap::new())); @@ -476,12 +477,9 @@ async fn main() -> Result<(), Box> { // One-shot mode if let Some(prompt) = args.prompt { - match engine.run(&prompt).await { - Ok(_) => {} - Err(e) => { - eprintln!("\n❌ Error: {}", e); - std::process::exit(1); - } + let output = engine.run(&mut history, &prompt).await; + if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = output.result { + eprintln!("\n❌ Error: {error}"); } return Ok(()); @@ -500,13 +498,8 @@ async fn main() -> Result<(), Box> { return Ok(()); } - let mut locked = match engine.run(first_input).await { - Ok(out) => out.engine, - Err(e) => { - eprintln!("\n❌ Error: {}", e); - return Ok(()); - } - }; + let output = engine.run(&mut history, first_input).await; + let mut locked = output.engine; loop { print!("\n👤 You: "); @@ -525,11 +518,10 @@ async fn main() -> Result<(), Box> { break; } - match locked.run(input).await { - Ok(_) => {} - Err(e) => { - eprintln!("\n❌ Error: {}", e); - } + if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = + locked.run(&mut history, input).await + { + eprintln!("\n❌ Error: {error}"); } } diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 6619997a..e9bb671b 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -1,13 +1,14 @@ -use std::collections::HashMap; -use std::{marker::PhantomData, sync::Arc, time::Instant}; +use std::collections::{HashMap, HashSet}; +use std::{future::Future, marker::PhantomData, pin::Pin, sync::Arc, time::Instant}; use futures::StreamExt; use serde_json::{Value, json}; use tokio::sync::mpsc; +use tokio::time::Instant as TokioInstant; use tracing::{debug, info, trace, warn}; use crate::{ - Item, + History, HistoryEntry, Item, callback::{ ClosureMetaHandler, ClosureTextBlockHandler, ClosureThinkingBlockHandler, ClosureToolUseBlockHandler, TextBlockScope, ThinkingBlockScope, ToolUseBlockScope, @@ -27,7 +28,8 @@ use crate::{ timeline::{TextBlockCollector, ThinkingBlockCollector, Timeline, ToolCallCollector}, tool::{ ToolCall, ToolDefinition as EngineToolDefinition, ToolError, ToolExecutionContext, - ToolOutputLimits, ToolResult, truncate_content, + ToolExecutionHandle, ToolExecutionPolicy, ToolExecutionTerminal, ToolOutputLimits, + ToolResult, ToolResultDisposition, truncate_content, }, tool_server::{ToolServer, ToolServerHandle}, }; @@ -47,12 +49,18 @@ pub enum EngineError { /// Cancelled by CancellationToken #[error("Cancelled")] Cancelled, + /// Paused by the caller at the next safe boundary. + #[error("Paused")] + PauseRequested, /// Config warnings (unsupported options) #[error("Config warnings: {}", .0.iter().map(|w| w.to_string()).collect::>().join(", "))] ConfigWarnings(Vec), /// A durable-history observer rejected an item before it entered history. #[error("History append failed: {0}")] HistoryAppend(String), + /// Tool terminalization lost its execution-attempt compare-and-set fence. + #[error("Tool execution attempt fence failed: {0}")] + ToolAttemptFence(String), } /// Tool registration error @@ -70,40 +78,178 @@ pub struct EngineConfig { _private: (), } -/// Engine execution result (status) +/// 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. +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; + 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, + } + } + + 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; + } + + ProviderHistoryProjection { + items: materialized.into_iter().map(|(_, item)| item).collect(), + original_to_projected_index, + } +} + +/// Legacy serializable outcome used by the Worker session-log compatibility boundary. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)] #[serde(rename_all = "snake_case")] pub enum EngineResult { - /// Completed (waiting for user input) Finished, - /// Paused (can be resumed) Paused, - /// Turn limit reached (max_turns exceeded) LimitReached, - /// Yielded to caller for external processing (e.g. context compaction). - /// - /// Distinct from `Paused`: internal machinery, not user-facing. The - /// caller is expected to perform some side work and then call `resume()` - /// to continue the turn loop. Yielded, } +/// The public termination boundary for one logical engine run. +#[derive(Debug)] +pub enum EngineRunExit { + Finished, + Paused, + Yielded, + Interrupted(StopReason), +} + +/// A typed reason why an engine run could not finish normally. +#[derive(Debug)] +pub enum StopReason { + LimitReached, + ContextWindowExceeded, + Cancelled, + Unexpected(EngineError), +} + +impl From> for EngineRunExit { + fn from(result: Result) -> Self { + match result { + Ok(EngineResult::Finished) => Self::Finished, + Ok(EngineResult::Paused) => Self::Paused, + Ok(EngineResult::Yielded) => Self::Yielded, + Ok(EngineResult::LimitReached) => Self::Interrupted(StopReason::LimitReached), + Err(EngineError::Client(ClientError::ContextWindowExceeded)) => { + Self::Interrupted(StopReason::ContextWindowExceeded) + } + Err(EngineError::Cancelled) => Self::Interrupted(StopReason::Cancelled), + Err(EngineError::PauseRequested) => Self::Paused, + Err(error) => Self::Interrupted(StopReason::Unexpected(error)), + } + } +} + /// Result of [`Engine::run`] or [`Engine::resume`]. /// /// Contains the `Locked` Engine (ready for subsequent runs) and the outcome. -pub struct EngineRunOutput { +pub struct EngineRunOutput { /// The Engine, now in Locked state. - pub engine: Engine, + pub engine: Engine, /// Outcome of the turn. - pub result: EngineResult, + pub result: EngineRunExit, } /// Internal: tool execution result enum ToolExecutionResult { - Completed(Vec), + Completed, Paused, } +#[derive(Debug, Clone)] +struct ToolExecutionAttempt { + attempt_id: String, + terminal: bool, +} + +/// Per-batch compare-and-set fence for terminal ToolResult commits. +/// +/// A completion may commit only when its attempt id still matches the active +/// execution for that call and no prior terminal output has won the fence. +#[derive(Debug, Default)] +struct ToolExecutionAttemptFence { + attempts: HashMap, +} + +impl ToolExecutionAttemptFence { + fn register(&mut self, call_id: String, attempt_id: String) { + self.attempts.insert( + call_id, + ToolExecutionAttempt { + attempt_id, + terminal: false, + }, + ); + } + + fn can_commit(&self, call_id: &str, attempt_id: &str) -> bool { + matches!( + self.attempts.get(call_id), + Some(attempt) if attempt.attempt_id == attempt_id && !attempt.terminal + ) + } + + fn commit_terminal(&mut self, call_id: &str, attempt_id: &str) -> bool { + let Some(attempt) = self.attempts.get_mut(call_id) else { + return false; + }; + if attempt.attempt_id != attempt_id || attempt.terminal { + return false; + } + attempt.terminal = true; + true + } + + fn is_terminal(&self, call_id: &str) -> bool { + self.attempts + .get(call_id) + .is_some_and(|attempt| attempt.terminal) + } + + #[cfg(test)] + fn attempt_id(&self, call_id: &str) -> Option<&str> { + self.attempts + .get(call_id) + .map(|attempt| attempt.attempt_id.as_str()) + } +} + const MAX_STREAM_CONTINUATIONS: u32 = 3; /// Central component for managing LLM interactions @@ -113,29 +259,31 @@ const MAX_STREAM_CONTINUATIONS: u32 = 3; /// /// # State Transitions (Type-state) /// -/// - [`Mutable`]: Initial state. System prompt, history, and tools can be freely edited. +/// - [`Mutable`]: Initial state. System prompt and tools can be edited; history is caller-owned. /// - [`Locked`]: Cache-protected state. Prefix context is immutable; only `run()` / `resume()` are available. /// /// Calling `run()` on a `Mutable` Engine consumes it and returns a -/// `Locked` Engine together with the result. This ensures the -/// cache prefix is fixed for optimal KV cache hit rate. +/// `Locked` Engine together with the result. The engine borrows the caller's +/// [`History`](crate::History) only while running, so host annotations stay with +/// the host-owned history and are never projected to providers. /// /// ```ignore +/// let mut history = History::new(); /// let mut engine = Engine::new(client) /// .system_prompt("You are a helpful assistant."); /// engine.register_tool(my_tool); /// /// // Mutable::run() consumes self → EngineRunOutput { engine: Locked, result } -/// let out = engine.run("Hello").await?; +/// let out = engine.run(&mut history, "Hello").await?; /// let mut engine = out.engine; /// /// // Locked::run() borrows &mut self -/// engine.run("Follow-up").await?; +/// engine.run(&mut history, "Follow-up").await?; /// /// // To edit between turns, unlock back to Mutable /// let mut engine = engine.unlock(); -/// engine.truncate_history(5); -/// let out = engine.run("Continue").await?; +/// history.truncate(5); +/// let out = engine.run(&mut history, "Continue").await?; /// let mut engine = out.engine; /// ``` #[derive(Debug, Clone, PartialEq, Eq)] @@ -155,7 +303,7 @@ enum StreamCompletion { Interrupted { reason: String }, } -pub struct Engine { +pub struct Engine { /// LLM client client: C, /// Retry policy for opening an LLM response stream. @@ -175,8 +323,6 @@ pub struct Engine { interceptor: Box, /// System prompt system_prompt: Option, - /// Item history (owned by Engine) - history: Vec, /// History length at lock time (only meaningful in Locked state) locked_prefix_len: usize, /// AgentTurn count across the lifetime of this Engine. @@ -202,6 +348,8 @@ pub struct Engine { tool_execution_batch_count: usize, /// Maximum number of AgentTurns (None = unlimited) max_turns: Option, + /// Caller-selected policy for interrupting started provider operations. + tool_execution_policy: ToolExecutionPolicy, /// AgentTurn-start callbacks (1:1 with LlmCall today) turn_start_cbs: Vec>, /// AgentTurn-end callbacks (1:1 with LlmCall today) @@ -237,11 +385,13 @@ pub struct Engine { history_append_cbs: Vec Result<(), String> + Send + Sync>>, /// Request configuration (max_tokens, temperature, etc.) request_config: RequestConfig, - /// Whether the previous run was interrupted - last_run_interrupted: bool, /// Cancel notification channel (for interrupting execution) cancel_tx: mpsc::Sender<()>, cancel_rx: mpsc::Receiver<()>, + /// Pause notification channel. Unlike cancellation, this waits for already + /// started tools to reach provider-confirmed terminal results. + pause_tx: mpsc::Sender<()>, + pause_rx: mpsc::Receiver<()>, /// Byte-size caps applied to tool `content` before it reaches history. /// `None` disables truncation (tests and minimal setups). tool_output_limits: Option, @@ -266,14 +416,10 @@ pub struct Engine { /// stable conversation identifier when the backend benefits from one. cache_key: Option, /// State marker - _state: PhantomData, + _state: PhantomData<(S, A)>, } -impl Engine { - fn reset_interruption_state(&mut self) { - self.last_run_interrupted = false; - } - +impl Engine { fn start_logical_run(&mut self) { self.active_run_turn_count = Some(0); } @@ -283,7 +429,10 @@ impl Engine { } fn finish_logical_run(&mut self, result: &Result) { - if !matches!(result, Ok(EngineResult::Paused) | Ok(EngineResult::Yielded)) { + if !matches!( + result, + Ok(EngineResult::Paused | EngineResult::Yielded) | Err(EngineError::PauseRequested) + ) { self.active_run_turn_count = None; } } @@ -292,14 +441,19 @@ impl Engine { while self.cancel_rx.try_recv().is_ok() {} } - /// Discard pending cancellation notifications while the engine is idle. + fn drain_pause_queue(&mut self) { + while self.pause_rx.try_recv().is_ok() {} + } + + /// Discard pending interruption notifications while the engine is idle. /// - /// Cancellation is a running-turn control signal. Callers that own a higher - /// level run state can use this before starting a new turn so an old idle - /// signal does not poison the next request, while cancellation queued after - /// the run has been accepted remains observable by the turn loop. + /// Cancellation and pause are running-turn control signals. Callers that own + /// a higher level run state can use this before starting a new turn so an old + /// idle signal does not poison the next request, while interruption queued + /// after the run has been accepted remains observable by the turn loop. pub fn clear_pending_cancel(&mut self) { self.drain_cancel_queue(); + self.drain_pause_queue(); } fn try_cancelled(&mut self) -> bool { @@ -311,6 +465,14 @@ impl Engine { } } + fn try_paused(&mut self) -> bool { + use tokio::sync::mpsc::error::TryRecvError; + match self.pause_rx.try_recv() { + Ok(()) => true, + Err(TryRecvError::Empty | TryRecvError::Disconnected) => false, + } + } + /// Register a text block observer with scoped callbacks. /// /// The setup closure is called once per text block. Inside it, register @@ -539,11 +701,15 @@ impl Engine { fn append_history_items( &mut self, + history: &mut History, items: impl IntoIterator, + annotate: &mut impl FnMut(&Item) -> Result, ) -> Result<(), EngineError> { for item in items { self.emit_history_append(&item)?; - self.history.push(item); + history + .append_with(item, annotate) + .map_err(EngineError::HistoryAppend)?; } Ok(()) } @@ -650,9 +816,9 @@ impl Engine { &self.client } - /// Get a reference to the history - pub fn history(&self) -> &[Item] { - &self.history + /// Borrow caller-owned annotated history entries. + pub fn history<'h>(&self, history: &'h History) -> &'h [HistoryEntry] { + history.entries() } /// Get a reference to the system prompt @@ -770,6 +936,22 @@ impl Engine { self.cancel_tx.clone() } + /// Get the safe-boundary pause notification sender. + pub fn pause_sender(&self) -> mpsc::Sender<()> { + self.pause_tx.clone() + } + + /// Select the deadline policy applied to already-started provider operations. + /// Worker/controller layers own this lifecycle policy; Agen owns only the + /// mechanical terminalization of each call. + pub fn set_tool_execution_policy(&mut self, policy: ToolExecutionPolicy) { + self.tool_execution_policy = policy; + } + + pub fn tool_execution_policy(&self) -> ToolExecutionPolicy { + self.tool_execution_policy + } + /// Set request configuration at once pub fn set_request_config(&mut self, config: RequestConfig) { self.request_config = config; @@ -805,11 +987,6 @@ impl Engine { self.try_cancelled() } - /// Whether the previous run was interrupted - pub fn last_run_interrupted(&self) -> bool { - self.last_run_interrupted - } - /// Generate list of ToolDefinitions for LLM from registered tools fn build_tool_definitions(&self) -> Vec { self.tool_server.tool_definitions_sorted() @@ -873,8 +1050,14 @@ impl Engine { request = request.system(system); } - // Add items directly (Request now uses Items natively) - request = request.items(context.iter().cloned()); + // History keeps terminal tool outputs in completion order so each + // result can be committed immediately. Providers, however, expect a + // deterministic projection matching the assistant's ToolCall order. + let projection = materialize_provider_history(context); + let projected_cache_anchor = self + .cache_anchor + .and_then(|anchor| projection.original_to_projected_index.get(anchor).copied()); + request = request.items(projection.items); // Add tool definitions for tool_def in tool_definitions { @@ -887,7 +1070,7 @@ impl Engine { // Attach the cache prefix anchor (may be narrower than `context` // if the prune projection trimmed items from the head — keep it // in range). - request.cache_anchor = self.cache_anchor.filter(|&anchor| anchor < context.len()); + request.cache_anchor = projected_cache_anchor; request.cache_key = self.cache_key.clone(); request @@ -902,7 +1085,6 @@ impl Engine { match result { Ok(value) => Ok(value), Err(err) => { - self.last_run_interrupted = true; let reason = match &err { EngineError::Aborted(reason) => reason.clone(), EngineError::Cancelled => "Cancelled".to_string(), @@ -915,20 +1097,20 @@ impl Engine { } /// Check for pending tool calls (for resuming from Pause) - fn get_pending_tool_calls(&self) -> Option> { + fn get_pending_tool_calls(&self, history: &History) -> Option> { // Find the last ToolCall items that don't have corresponding ToolResult let mut pending_calls = Vec::new(); let mut answered_call_ids = std::collections::HashSet::new(); // First pass: collect all answered call IDs - for item in &self.history { + 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 &self.history { + for item in history.items() { if let Item::ToolCall { call_id, name, @@ -960,9 +1142,17 @@ impl Engine { /// executes approved tools in parallel and applies post_tool_call hooks to results. async fn execute_tools( &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, tool_calls: Vec, ) -> Result { - use futures::future::join_all; + use futures::stream::{FuturesUnordered, StreamExt}; + + // A pause observed before provider ownership starts leaves every call + // NotStarted and therefore eligible for an explicit later retry. + if self.try_paused() { + return Ok(ToolExecutionResult::Paused); + } // Map from tool call ID to (ToolCall, Meta, Tool, Context) // Retained because it's needed for PostToolCall hooks @@ -1000,11 +1190,9 @@ impl Engine { continue; } PreToolAction::Abort(reason) => { - self.last_run_interrupted = true; return Err(EngineError::Aborted(reason)); } PreToolAction::Pause => { - self.last_run_interrupted = true; return Ok(ToolExecutionResult::Paused); } } @@ -1023,129 +1211,363 @@ impl Engine { context.clone(), ), ); - approved_calls.push((tool_call, context)); + approved_calls.push((tool_call, context, Some(info.tool))); } else { // Unknown tools go into approved list as-is (will error at execution) let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index); - approved_calls.push((tool_call, context)); + approved_calls.push((tool_call, context, None)); } } - // Phase 2: Execute approved tools in parallel (cancellable) - let futures: Vec<_> = approved_calls - .into_iter() - .map(|(tool_call, context)| { - let tool_server = self.tool_server.clone(); - async move { - let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default(); - match tool_server - .call_tool(&tool_call.name, &input_json, context) - .await - { - Ok(output) => ToolResult::from_output(&tool_call.id, output), - Err(e) => ToolResult::error(&tool_call.id, e.to_string()), - } - } - }) + // Phase 2: Execute approved tools in parallel. FuturesUnordered yields + // each terminal result as soon as that call completes instead of + // holding fast siblings behind the slowest call in the batch. + let started_calls: Vec<_> = approved_calls + .iter() + .map(|(tool_call, context, _)| (tool_call.id.clone(), context.batch_id.clone())) .collect(); - - // Make tool execution cancellable - let mut results = tokio::select! { - results = join_all(futures) => results, - cancel = self.cancel_rx.recv() => { - if cancel.is_some() { - info!("Tool execution cancelled"); + let mut attempt_fence = ToolExecutionAttemptFence::default(); + for (call_id, attempt_id) in &started_calls { + attempt_fence.register(call_id.clone(), attempt_id.clone()); + } + let futures: FuturesUnordered + Send>>> = + FuturesUnordered::new(); + let mut execution_handles = HashMap::new(); + for (tool_call, context, tool) in approved_calls { + let attempt_id = context.batch_id.clone(); + let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default(); + let call_id = tool_call.id.clone(); + let future: Pin + Send>> = match tool { + None => { + let result = + ToolResult::error(&call_id, format!("Tool not found: {}", tool_call.name)); + Box::pin(async move { (attempt_id, result) }) } - self.timeline.abort_current_block(); - self.last_run_interrupted = true; - return Err(EngineError::Cancelled); - } - }; - results.extend(synthetic_results); + Some(tool) => { + let (handle, terminal) = ToolExecutionHandle::start(tool, input_json, context); + execution_handles.insert(call_id.clone(), handle); + 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) + }) + } + }; + futures.push(future); + } - // Phase 3: Apply post_tool_call interceptor - for tool_result in &mut results { - if let Some((tool_call, meta, tool, context)) = - call_info_map.get(&tool_result.tool_use_id) - { - let mut info = ToolResultInfo { - call: tool_call.clone(), - result: tool_result.clone(), - meta: meta.clone(), - tool: tool.clone(), - context: context.clone(), - }; + // Synthetic results are already terminal and need no execution wait. + // Commit them before polling ordinary calls so they obey the same + // commit-before-publish boundary. + let mut terminal_call_ids = HashSet::new(); + let mut pause_requested = false; + let mut pause_deadline = None; + for result in synthetic_results { + self.finalize_and_commit_tool_result( + history, + annotate, + result, + None, + &call_info_map, + &mut attempt_fence, + &mut terminal_call_ids, + ) + .await?; + } - match self.interceptor.post_tool_call(&mut info).await { - PostToolAction::Continue => {} - PostToolAction::Abort(reason) => { - self.last_run_interrupted = true; - return Err(EngineError::Aborted(reason)); + let mut futures = futures; + while !futures.is_empty() { + tokio::select! { + // If cancellation and a completed result are both ready, drain + // the completed result first. This preserves every terminal + // output observed before the cancellation boundary. + biased; + result = futures.next() => { + let (attempt_id, result) = + result.expect("non-empty FuturesUnordered returns a result"); + self.finalize_and_commit_tool_result( + history, + annotate, + result, + Some(&attempt_id), + &call_info_map, + &mut attempt_fence, + &mut terminal_call_ids, + ).await?; + } + pause = self.pause_rx.recv(), if !pause_requested => { + if pause.is_some() { + // Pause first waits for already-started tools to reach a + // natural safe boundary. If they do not, Worker policy + // escalates to the same explicit cancel-and-confirm path. + pause_requested = true; + pause_deadline = Some( + TokioInstant::now() + + self.tool_execution_policy.pause_safe_boundary_timeout, + ); } } - // Reflect interceptor-modified results - *tool_result = info.result; - } - } + _ = tokio::time::sleep_until(pause_deadline.unwrap_or_else(TokioInstant::now)), if pause_deadline.is_some() => { + pause_deadline = None; + let _ = self.cancel_tx.try_send(()); + } + cancel = self.cancel_rx.recv() => { + if cancel.is_some() { + info!("Tool execution cancellation requested"); + } - // Phase 4: Cap `content` byte-size before it enters history. - // Runs *after* post_tool_call so interceptors (audit, logging, - // classification) still observe the full content, and any - // content they inject is also truncated — closing the last gap - // before the data reaches the next LLM request. - if let Some(limits) = self.tool_output_limits.as_ref() { - for tool_result in &mut results { - let Some(content) = tool_result.content.as_mut() else { - continue; - }; - let Some((tool_call, _, _, _)) = call_info_map.get(&tool_result.tool_use_id) else { - continue; - }; - let limit = limits.limit_for(&tool_call.name); - let before = content.len(); - truncate_content(content, limit); - if content.len() != before { - warn!( - tool = %tool_call.name, - before_bytes = before, - after_bytes = content.len(), - limit_bytes = limit, - "Tool output exceeded byte limit and was truncated" - ); - self.emit_warning(&format!( - "tool `{}` output truncated from {} to {} bytes (limit {})", - tool_call.name, - before, - content.len(), - limit - )); + let cancellation_request_deadline = TokioInstant::now() + + self.tool_execution_policy.cancellation_request_timeout; + let cancellation_requests = execution_handles + .iter() + .filter(|(call_id, _)| !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" + ); + } + } + + // Keep polling the original execution handles to their + // provider-confirmed terminal result until the caller-selected + // deadline. The execution remains owned even if this polling + // future is later dropped. + let deadline = TokioInstant::now() + + self.tool_execution_policy.terminal_confirmation_timeout; + while !futures.is_empty() { + tokio::select! { + biased; + result = futures.next() => { + let (attempt_id, result) = + result.expect("non-empty FuturesUnordered returns a result"); + self.finalize_and_commit_tool_result( + history, + annotate, + result, + Some(&attempt_id), + &call_info_map, + &mut attempt_fence, + &mut terminal_call_ids, + ).await?; + } + _ = tokio::time::sleep_until(deadline) => break, + } + } + + // Calls that did not confirm a terminal outcome inside the + // grace period are durably closed as OutcomeUnknown before + // Engine/Worker final status becomes observable. + for (call_id, attempt_id) in &started_calls { + if !attempt_fence.is_terminal(call_id) { + if let Some(handle) = execution_handles.get(call_id) { + handle.force_close(); + } + self.finalize_and_commit_tool_result( + history, + annotate, + ToolResult::outcome_unknown(call_id), + Some(attempt_id), + &call_info_map, + &mut attempt_fence, + &mut terminal_call_ids, + ).await?; + } + } + + self.timeline.abort_current_block(); + if pause_requested { + return Ok(ToolExecutionResult::Paused); + } + return Err(EngineError::Cancelled); } } } - // Emit per-result callbacks on the post-truncation payload. - for tool_result in &results { - self.emit_tool_result(tool_result); + Ok(if pause_requested { + ToolExecutionResult::Paused + } else { + ToolExecutionResult::Completed + }) + } + + /// Apply post-execution policy, bound the model-visible payload, durably + /// append one terminal ToolResult, and only then publish it to observers. + async fn finalize_and_commit_tool_result( + &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + mut tool_result: ToolResult, + execution_attempt_id: Option<&str>, + call_info_map: &HashMap< + String, + ( + ToolCall, + crate::tool::ToolMeta, + Arc, + ToolExecutionContext, + ), + >, + attempt_fence: &mut ToolExecutionAttemptFence, + terminal_call_ids: &mut HashSet, + ) -> Result { + let call_id = tool_result.tool_use_id.as_str(); + let may_commit = match execution_attempt_id { + Some(attempt_id) => attempt_fence.can_commit(call_id, attempt_id), + None => !terminal_call_ids.contains(call_id), + }; + if !may_commit { + warn!( + call_id, + execution_attempt_id, + disposition = ?tool_result.disposition, + "Ignoring stale or duplicate tool result after terminal output commit" + ); + return Ok(false); } - Ok(ToolExecutionResult::Completed(results)) + let call_info = call_info_map.get(&tool_result.tool_use_id); + let mut abort_reason = None; + if let Some((tool_call, meta, tool, context)) = call_info { + let mut info = ToolResultInfo { + call: tool_call.clone(), + result: tool_result, + meta: meta.clone(), + tool: tool.clone(), + context: context.clone(), + }; + + match self.interceptor.post_tool_call(&mut info).await { + PostToolAction::Continue => {} + PostToolAction::Abort(reason) => { + abort_reason = Some(reason); + } + } + tool_result = info.result; + } + if tool_result.is_error && tool_result.disposition.is_success() { + tool_result.disposition = ToolResultDisposition::Error; + } + tool_result.is_error = !tool_result.disposition.is_success(); + + // Cap content only after post_tool_call so interceptors still observe + // the full payload and any content they inject is bounded too. + if let (Some(limits), Some((tool_call, _, _, _)), Some(content)) = ( + self.tool_output_limits.as_ref(), + call_info, + tool_result.content.as_mut(), + ) { + let limit = limits.limit_for(&tool_call.name); + let before = content.len(); + truncate_content(content, limit); + if content.len() != before { + warn!( + tool = %tool_call.name, + before_bytes = before, + after_bytes = content.len(), + limit_bytes = limit, + "Tool output exceeded byte limit and was truncated" + ); + self.emit_warning(&format!( + "tool `{}` output truncated from {} to {} bytes (limit {})", + tool_call.name, + before, + content.len(), + limit + )); + } + } + + let item = Item::tool_result_item_with_disposition_and_attachments( + &tool_result.tool_use_id, + &tool_result.summary, + tool_result.content.clone(), + tool_result.disposition, + tool_result.attachments.clone(), + ); + self.append_history_items(history, std::iter::once(item), annotate)?; + if let Some(attempt_id) = execution_attempt_id + && !attempt_fence.commit_terminal(&tool_result.tool_use_id, attempt_id) + { + return Err(EngineError::ToolAttemptFence( + "tool execution attempt fence changed during terminal commit".to_string(), + )); + } + terminal_call_ids.insert(tool_result.tool_use_id.clone()); + debug!( + tool = call_info + .map(|(call, _, _, _)| call.name.as_str()) + .unwrap_or("unknown"), + call_id = %tool_result.tool_use_id, + execution_attempt_id, + disposition = ?tool_result.disposition, + "Tool execution terminalized" + ); + self.emit_tool_result(&tool_result); + if let Some(reason) = abort_reason { + return Err(EngineError::Aborted(reason)); + } + Ok(true) } /// Internal turn execution logic - async fn run_turn_loop(&mut self) -> Result { - self.reset_interruption_state(); + async fn run_turn_loop( + &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + ) -> Result { let tool_definitions = self.build_tool_definitions(); info!( - item_count = self.history.len(), + item_count = history.len(), tool_count = tool_definitions.len(), "Starting engine run" ); // Resume pending tool calls from a previous Pause - if let Some(tool_calls) = self.get_pending_tool_calls() { + if let Some(tool_calls) = self.get_pending_tool_calls(history) { info!("Resuming pending tool calls"); - if let Some(result) = self.execute_and_commit_tools(tool_calls).await? { + if let Some(result) = self + .execute_and_commit_tools(history, annotate, tool_calls) + .await? + { return Ok(result); } } @@ -1156,7 +1578,6 @@ impl Engine { if self.try_cancelled() { info!("Execution cancelled"); self.timeline.abort_current_block(); - self.last_run_interrupted = true; return Err(EngineError::Cancelled); } @@ -1169,7 +1590,6 @@ impl Engine { max_turns = max, "Logical run turn limit reached" ); - self.last_run_interrupted = false; return Ok(EngineResult::LimitReached); } @@ -1192,13 +1612,13 @@ impl Engine { .await .map_err(EngineError::HistoryAppend)?; if !pending.is_empty() { - self.append_history_items(pending)?; + self.append_history_items(history, pending, annotate)?; } // Clone the history into a per-request context. Everything // below (prune projection, interceptor hooks) mutates only - // this clone, so the persistent `self.history` stays intact. - let mut request_context = self.history.clone(); + // this clone, so the caller-owned `history` stays intact. + let mut request_context = history.items_cloned(); // Prune projection: if both the config and the savings // estimator are configured, drop ToolResult.content from @@ -1263,17 +1683,15 @@ impl Engine { for cb in &self.turn_end_cbs { cb(current_turn); } - self.last_run_interrupted = true; return Err(EngineError::Aborted(reason)); } PreRequestAction::YieldWith(items) => { - self.append_history_items(items.clone())?; + self.append_history_items(history, items.clone(), annotate)?; request_context.extend(items); info!("Yielded by interceptor after pre-request history append"); for cb in &self.turn_end_cbs { cb(current_turn); } - self.last_run_interrupted = true; return Ok(EngineResult::Yielded); } PreRequestAction::Yield => { @@ -1281,11 +1699,10 @@ impl Engine { for cb in &self.turn_end_cbs { cb(current_turn); } - self.last_run_interrupted = true; return Ok(EngineResult::Yielded); } PreRequestAction::ContinueWith(items) => { - self.append_history_items(items.clone())?; + self.append_history_items(history, items.clone(), annotate)?; request_context.extend(items); } PreRequestAction::Continue => {} @@ -1326,7 +1743,6 @@ impl Engine { if let StreamCompletion::Interrupted { reason } = stream_outcome { stream_continuations += 1; if stream_continuations > MAX_STREAM_CONTINUATIONS { - self.last_run_interrupted = true; return Err(EngineError::Client(ClientError::Api { status: None, code: None, @@ -1345,7 +1761,7 @@ impl Engine { let assistant_items = self.build_assistant_items(&reasoning_items, &text_blocks, &[]); if !assistant_items.is_empty() { - self.append_history_items(assistant_items)?; + self.append_history_items(history, assistant_items, annotate)?; } self.emit_llm_continuation( current_llm_call, @@ -1373,26 +1789,28 @@ impl Engine { let tool_calls = self.tool_call_collector.take_collected(); let assistant_items = self.build_assistant_items(&reasoning_items, &text_blocks, &tool_calls); - self.append_history_items(assistant_items)?; + self.append_history_items(history, assistant_items, annotate)?; if tool_calls.is_empty() { - match self.interceptor.on_turn_end(&self.history).await { + let turn_end_context = history.items_cloned(); + match self.interceptor.on_turn_end(&turn_end_context).await { TurnEndAction::Finish => { - self.last_run_interrupted = false; return Ok(EngineResult::Finished); } TurnEndAction::ContinueWithMessages(additional) => { - self.append_history_items(additional)?; + self.append_history_items(history, additional, annotate)?; continue; } TurnEndAction::Pause => { - self.last_run_interrupted = true; return Ok(EngineResult::Paused); } } } - if let Some(result) = self.execute_and_commit_tools(tool_calls).await? { + if let Some(result) = self + .execute_and_commit_tools(history, annotate, tool_calls) + .await? + { return Ok(result); } } @@ -1422,6 +1840,13 @@ impl Engine { let stream_started = Instant::now(); let stream_result = tokio::select! { stream_result = self.client.stream(request.clone()) => stream_result, + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!("Paused before stream started"); + } + self.timeline.abort_current_block(); + return Err(EngineError::PauseRequested); + } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!("Cancelled before stream started"); @@ -1436,7 +1861,6 @@ impl Engine { }), ); self.timeline.abort_current_block(); - self.last_run_interrupted = true; return Err(EngineError::Cancelled); } }; @@ -1454,6 +1878,13 @@ impl Engine { ); let first_event_result = tokio::select! { first_event = wait_for_first_stream_event(stream, DEFAULT_FIRST_STREAM_EVENT_TIMEOUT) => first_event, + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!("Paused before first stream event"); + } + self.timeline.abort_current_block(); + return Err(EngineError::PauseRequested); + } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!("Cancelled before first stream event"); @@ -1468,7 +1899,6 @@ impl Engine { }), ); self.timeline.abort_current_block(); - self.last_run_interrupted = true; return Err(EngineError::Cancelled); } }; @@ -1510,7 +1940,6 @@ impl Engine { let next_failed_attempt = failed_attempt + 1; if next_failed_attempt >= policy.max_attempts || !is_retryable(&err) { - self.last_run_interrupted = true; return Err(EngineError::Client(err)); } @@ -1519,7 +1948,6 @@ impl Engine { .unwrap_or_else(|| policy.backoff(failed_attempt)); let elapsed = started.elapsed(); if elapsed + wait > policy.total_timeout { - self.last_run_interrupted = true; return Err(EngineError::Client(err)); } @@ -1543,12 +1971,18 @@ impl Engine { tokio::select! { _ = tokio::time::sleep(wait) => {} + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!("Paused during LLM retry backoff"); + } + self.timeline.abort_current_block(); + return Err(EngineError::PauseRequested); + } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!("Cancelled during LLM retry backoff"); } self.timeline.abort_current_block(); - self.last_run_interrupted = true; return Err(EngineError::Cancelled); } } @@ -1591,7 +2025,6 @@ impl Engine { let event = match result { Ok(event) => event, Err(err) => { - self.last_run_interrupted = true; // 部分情報でも発火しておく(料金会計用) self.timeline.flush_usage(); return Ok(StreamCompletion::Interrupted { @@ -1612,7 +2045,6 @@ impl Engine { if let Event::Error(err) = &event { self.timeline.abort_current_block(); self.timeline.flush_usage(); - self.last_run_interrupted = true; return Err(EngineError::Client(ClientError::Api { status: None, code: err.code.clone(), @@ -1624,13 +2056,19 @@ impl Engine { None => break, } } + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!("Paused during response stream"); + } + self.timeline.abort_current_block(); + return Err(EngineError::PauseRequested); + } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!("Stream cancelled"); } self.timeline.abort_current_block(); self.timeline.flush_usage(); - self.last_run_interrupted = true; return Err(EngineError::Cancelled); } } @@ -1646,44 +2084,27 @@ impl Engine { /// `None` if the turn loop should continue. async fn execute_and_commit_tools( &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, tool_calls: Vec, ) -> Result, EngineError> { - match self.execute_tools(tool_calls).await { - Ok(ToolExecutionResult::Paused) => { - self.last_run_interrupted = true; - Ok(Some(EngineResult::Paused)) - } - Ok(ToolExecutionResult::Completed(results)) => { - // Route per-result pushes through the callback path so - // observers see each tool result as it lands. - let items = results.into_iter().map(|result| { - Item::tool_result_item_with_attachments( - &result.tool_use_id, - &result.summary, - result.content, - result.is_error, - result.attachments, - ) - }); - self.append_history_items(items)?; - Ok(None) - } - Err(err) => { - self.last_run_interrupted = true; - Err(err) - } + match self.execute_tools(history, annotate, tool_calls).await { + Ok(ToolExecutionResult::Paused) => Ok(Some(EngineResult::Paused)), + Ok(ToolExecutionResult::Completed) => Ok(None), + Err(err) => Err(err), } } } -impl Engine { - /// Create a new Engine (in Mutable state) - pub fn new(client: C) -> Self { +impl Engine { + /// Create a new annotated Engine (in Mutable state). + pub fn new_annotated(client: C) -> Self { let text_block_collector = TextBlockCollector::new(); let tool_call_collector = ToolCallCollector::new(); let thinking_block_collector = ThinkingBlockCollector::new(); let mut timeline = Timeline::new(); let (cancel_tx, cancel_rx) = mpsc::channel(1); + let (pause_tx, pause_rx) = mpsc::channel(1); // Register collectors with Timeline timeline.on_text_block(text_block_collector.clone()); @@ -1700,13 +2121,13 @@ impl Engine { tool_server: ToolServer::new().handle(), interceptor: Box::new(DefaultInterceptor), system_prompt: None, - history: Vec::new(), locked_prefix_len: 0, turn_count: 0, active_run_turn_count: None, llm_call_count: 0, tool_execution_batch_count: 0, max_turns: None, + tool_execution_policy: ToolExecutionPolicy::default(), turn_start_cbs: Vec::new(), turn_end_cbs: Vec::new(), llm_call_start_cbs: Vec::new(), @@ -1719,9 +2140,10 @@ impl Engine { tool_result_cbs: Vec::new(), history_append_cbs: Vec::new(), request_config: RequestConfig::default(), - last_run_interrupted: false, cancel_tx, cancel_rx, + pause_tx, + pause_rx, tool_output_limits: None, prune_config: None, token_estimator: None, @@ -1861,36 +2283,38 @@ impl Engine { } } - /// Replace history during restore/rebuild without emitting append callbacks. + /// Replace caller-owned history during restore/rebuild without emitting append callbacks. /// /// This is not a history-growth API. Live append paths must use - /// [`append_history`](Self::append_history) so `on_history_append` observers - /// see every inserted item. - pub fn set_history(&mut self, items: Vec) { - self.history = items; - } - - /// Append items to history after every history-append observer accepts the - /// item. This is the only public Mutable-state API for growing engine - /// history; callers that need session-log persistence must install - /// [`on_history_append`](Self::on_history_append) before calling it. - pub fn append_history( + /// [`append_history_with`](Self::append_history_with) so observers and the + /// trusted annotation callback see every inserted item. + pub fn replace_history_entries( &mut self, + history: &mut History, + entries: Vec>, + ) -> Vec> { + history.replace_entries(entries) + } + + /// Append items to caller-owned history after every observer and the trusted + /// annotation callback accepts the item. + pub fn append_history_with( + &mut self, + history: &mut History, items: impl IntoIterator, + annotate: &mut impl FnMut(&Item) -> Result, ) -> Result<(), EngineError> { - self.append_history_items(items) + self.append_history_items(history, items, annotate) } - /// Truncate history without emitting append callbacks. - /// - /// This is an edit operation, not a history-growth path. - pub fn truncate_history(&mut self, len: usize) { - self.history.truncate(len); + /// Truncate caller-owned history without emitting append callbacks. + pub fn truncate_history(&mut self, history: &mut History, len: usize) { + history.truncate(len); } - /// Clear history - pub fn clear_history(&mut self) { - self.history.clear(); + /// Clear caller-owned history. + pub fn clear_history(&mut self, history: &mut History) { + history.clear(); } /// Set the turn count (for session restoration) @@ -1903,49 +2327,47 @@ impl Engine { self.max_turns = max_turns; } - /// Set the last_run_interrupted flag (for session restoration) - pub fn set_last_run_interrupted(&mut self, interrupted: bool) { - self.last_run_interrupted = interrupted; - if !interrupted { - self.active_run_turn_count = None; - } - } - /// Apply configuration (reserved for future extensions) #[allow(dead_code)] pub fn config(self, _config: EngineConfig) -> Self { self } - /// Execute a turn, consuming self and transitioning to Locked. + /// Run the engine with one user input, appending to caller-owned history. /// - /// This is the primary entry point for first use. Equivalent to - /// `self.lock()` followed by `locked.run(user_input)`. - /// - /// Subsequent runs can call [`Engine::run`] directly. - /// To edit state between turns, call [`unlock()`](Engine::unlock) first. - pub async fn run( + /// The trusted `annotate` callback is invoked after append observers and before + /// each new item becomes live in `history`. Providers, token counters, pruners, + /// and interceptors receive only the `Item` projection. + pub async fn run_with_annotation( self, + history: &mut History, user_input: impl Into, - ) -> Result, EngineError> { - let mut locked = self.lock(); - let result = locked.run(user_input).await?; - Ok(EngineRunOutput { + annotate: &mut impl FnMut(&Item) -> Result, + ) -> EngineRunOutput { + let mut locked = self.lock(history); + let result = locked + .run_with_annotation(history, user_input, annotate) + .await; + EngineRunOutput { engine: locked, result, - }) + } } /// Resume from Paused, consuming self and transitioning to Locked. /// /// Used after `unlock()` → edit → resume. - pub async fn resume(self) -> Result, EngineError> { - let mut locked = self.lock(); - let result = locked.resume().await?; - Ok(EngineRunOutput { + pub async fn resume_with_annotation( + self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + ) -> EngineRunOutput { + let mut locked = self.lock(history); + let result = locked.resume_with_annotation(history, annotate).await; + EngineRunOutput { engine: locked, result, - }) + } } /// Lock and transition to Locked state @@ -1961,9 +2383,9 @@ impl Engine { /// # Panics /// /// Panics if a pending tool factory produces a duplicate name. - pub fn lock(self) -> Engine { + pub fn lock(self, history: &History) -> Engine { self.tool_server.flush_pending(); - let locked_prefix_len = self.history.len(); + let locked_prefix_len = history.len(); Engine { client: self.client, retry_policy: self.retry_policy, @@ -1974,13 +2396,13 @@ impl Engine { tool_server: self.tool_server, interceptor: self.interceptor, system_prompt: self.system_prompt, - history: self.history, locked_prefix_len, turn_count: self.turn_count, active_run_turn_count: self.active_run_turn_count, llm_call_count: self.llm_call_count, tool_execution_batch_count: self.tool_execution_batch_count, max_turns: self.max_turns, + tool_execution_policy: self.tool_execution_policy, turn_start_cbs: self.turn_start_cbs, turn_end_cbs: self.turn_end_cbs, llm_call_start_cbs: self.llm_call_start_cbs, @@ -1993,10 +2415,11 @@ impl Engine { tool_result_cbs: self.tool_result_cbs, history_append_cbs: self.history_append_cbs, request_config: self.request_config, - last_run_interrupted: self.last_run_interrupted, cancel_tx: self.cancel_tx, cancel_rx: self.cancel_rx, + pause_tx: self.pause_tx, + pause_rx: self.pause_rx, tool_output_limits: self.tool_output_limits, prune_config: self.prune_config, token_estimator: self.token_estimator, @@ -2009,23 +2432,76 @@ impl Engine { } } -impl Engine { +fn unit_history_annotation(_: &Item) -> Result<(), String> { + Ok(()) +} + +impl Engine { + /// Create a new Engine (in Mutable state) using unit history annotations. + pub fn new(client: C) -> Self { + Self::new_annotated(client) + } + + /// Append unit-annotated items to caller-owned history. + pub fn append_history( + &mut self, + history: &mut History<()>, + items: impl IntoIterator, + ) -> Result<(), EngineError> { + let mut annotate = unit_history_annotation; + self.append_history_items(history, items, &mut annotate) + } + + /// Replace unit-annotated history from plain items. + pub fn set_history(&mut self, history: &mut History<()>, items: Vec) { + history.replace_items(items); + } + + /// Run using unit annotations. + pub async fn run( + self, + history: &mut History<()>, + user_input: impl Into, + ) -> EngineRunOutput { + let mut annotate = unit_history_annotation; + self.run_with_annotation(history, user_input, &mut annotate) + .await + } + + /// Resume using unit annotations. + pub async fn resume(self, history: &mut History<()>) -> EngineRunOutput { + let mut annotate = unit_history_annotation; + self.resume_with_annotation(history, &mut annotate).await + } +} + +impl Engine { /// Execute a turn /// /// Adds a new user message to history and sends a request to the LLM. /// Automatically loops if there are tool calls. - pub async fn run( + pub async fn run_with_annotation( &mut self, + history: &mut History, user_input: impl Into, + annotate: &mut impl FnMut(&Item) -> Result, + ) -> EngineRunExit { + self.run_result_with_annotation(history, user_input.into(), annotate) + .await + .into() + } + + async fn run_result_with_annotation( + &mut self, + history: &mut History, + user_input: String, + annotate: &mut impl FnMut(&Item) -> Result, ) -> Result { // Supplying new user input abandons any paused/yielded logical run. self.active_run_turn_count = None; - self.reset_interruption_state(); - // Interceptor: on_prompt_submit let mut user_item = Item::user_message(user_input); let extras = match self.interceptor.on_prompt_submit(&mut user_item).await { PromptAction::Cancel(reason) => { - self.last_run_interrupted = true; return self .finalize_interruption(Err(EngineError::Aborted(reason))) .await; @@ -2033,24 +2509,41 @@ impl Engine { PromptAction::Continue => Vec::new(), PromptAction::ContinueWith(items) => items, }; - self.append_history_items(std::iter::once(user_item))?; + self.append_history_items(history, std::iter::once(user_item), annotate)?; if !extras.is_empty() { - self.append_history_items(extras)?; + self.append_history_items(history, extras, annotate)?; } self.start_logical_run(); - let result = self.run_turn_loop().await; + let result = match self.run_turn_loop(history, annotate).await { + Err(EngineError::PauseRequested) => Ok(EngineResult::Paused), + other => other, + }; let result = self.finalize_interruption(result).await; self.finish_logical_run(&result); result } - /// Resume execution (from Paused state) - /// - /// Resumes turn processing from current state without adding a new user message. - pub async fn resume(&mut self) -> Result { - self.reset_interruption_state(); + /// Resume execution (from Paused state). + pub async fn resume_with_annotation( + &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + ) -> EngineRunExit { + self.resume_result_with_annotation(history, annotate) + .await + .into() + } + + async fn resume_result_with_annotation( + &mut self, + history: &mut History, + annotate: &mut impl FnMut(&Item) -> Result, + ) -> Result { self.ensure_logical_run(); - let result = self.run_turn_loop().await; + let result = match self.run_turn_loop(history, annotate).await { + Err(EngineError::PauseRequested) => Ok(EngineResult::Paused), + other => other, + }; let result = self.finalize_interruption(result).await; self.finish_logical_run(&result); result @@ -2065,7 +2558,7 @@ impl Engine { /// /// Note: After this operation, subsequent requests may not hit the cache. /// Use only when you need to edit history. - pub fn unlock(self) -> Engine { + pub fn unlock(self) -> Engine { Engine { client: self.client, retry_policy: self.retry_policy, @@ -2076,13 +2569,13 @@ impl Engine { tool_server: self.tool_server, interceptor: self.interceptor, system_prompt: self.system_prompt, - history: self.history, locked_prefix_len: 0, turn_count: self.turn_count, active_run_turn_count: self.active_run_turn_count, llm_call_count: self.llm_call_count, tool_execution_batch_count: self.tool_execution_batch_count, max_turns: self.max_turns, + tool_execution_policy: self.tool_execution_policy, turn_start_cbs: self.turn_start_cbs, turn_end_cbs: self.turn_end_cbs, llm_call_start_cbs: self.llm_call_start_cbs, @@ -2095,10 +2588,11 @@ impl Engine { tool_result_cbs: self.tool_result_cbs, history_append_cbs: self.history_append_cbs, request_config: self.request_config, - last_run_interrupted: self.last_run_interrupted, cancel_tx: self.cancel_tx, cancel_rx: self.cancel_rx, + pause_tx: self.pause_tx, + pause_rx: self.pause_rx, tool_output_limits: self.tool_output_limits, prune_config: self.prune_config, token_estimator: self.token_estimator, @@ -2111,6 +2605,25 @@ impl Engine { } } +impl Engine { + /// Run another turn using unit annotations. + pub async fn run( + &mut self, + history: &mut History<()>, + user_input: impl Into, + ) -> EngineRunExit { + let mut annotate = unit_history_annotation; + self.run_with_annotation(history, user_input, &mut annotate) + .await + } + + /// Resume using unit annotations. + pub async fn resume(&mut self, history: &mut History<()>) -> EngineRunExit { + let mut annotate = unit_history_annotation; + self.resume_with_annotation(history, &mut annotate).await + } +} + enum FirstStreamEvent { Ready(ResponseStream), Empty(ResponseStream), @@ -2215,6 +2728,43 @@ mod tests { use crate::tool::{Attachment, ImageAttachment}; use std::time::Duration; + #[test] + fn tool_execution_attempt_fence_rejects_duplicate_and_stale_results() { + let mut fence = ToolExecutionAttemptFence::default(); + fence.register("call".to_string(), "attempt-1".to_string()); + assert!(fence.can_commit("call", "attempt-1")); + assert!(fence.commit_terminal("call", "attempt-1")); + assert!(!fence.commit_terminal("call", "attempt-1")); + + fence.register("call".to_string(), "attempt-2".to_string()); + assert_eq!(fence.attempt_id("call"), Some("attempt-2")); + assert!(!fence.can_commit("call", "attempt-1")); + assert!(!fence.commit_terminal("call", "attempt-1")); + assert!(fence.commit_terminal("call", "attempt-2")); + } + + #[test] + fn provider_projection_reorders_results_and_remaps_cache_anchor() { + let items = vec![ + Item::tool_call_json("call_slow", "slow", serde_json::json!({})), + Item::tool_call_json("call_fast", "fast", serde_json::json!({})), + Item::tool_result_item("call_fast", "fast result", None, false), + Item::tool_result_item("call_slow", "slow result", None, false), + ]; + + let projection = materialize_provider_history(&items); + let result_order: Vec<_> = projection + .items + .iter() + .filter_map(|item| match item { + Item::ToolResult { call_id, .. } => Some(call_id.as_str()), + _ => None, + }) + .collect(); + assert_eq!(result_order, ["call_slow", "call_fast"]); + assert_eq!(projection.original_to_projected_index, [0, 1, 3, 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/history.rs b/crates/agen/src/history.rs new file mode 100644 index 00000000..4bb198cb --- /dev/null +++ b/crates/agen/src/history.rs @@ -0,0 +1,199 @@ +//! Typed conversation history containers. +//! +//! Agen keeps provider-visible [`Item`](crate::Item) values separate from any +//! host-domain provenance. The host chooses the annotation type `A`, while Agen +//! preserves each item and annotation as one entry for clone/truncate/restore +//! style history operations. + +use serde::{Deserialize, Serialize}; + +use crate::Item; + +/// One conversation-history entry with host-owned annotation. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct HistoryEntry { + /// Provider/model-visible conversation item. + pub item: Item, + /// Host-domain metadata kept with the item and never projected to providers. + pub annotation: A, +} + +impl HistoryEntry { + /// Build an entry from an item and its annotation. + pub fn new(item: Item, annotation: A) -> Self { + Self { item, annotation } + } + + /// Split the entry into its item and annotation. + pub fn into_parts(self) -> (Item, A) { + (self.item, self.annotation) + } +} + +impl HistoryEntry<()> { + /// Build a unit-annotated entry. + pub fn from_item(item: Item) -> Self { + Self { + item, + annotation: (), + } + } +} + +/// Conversation history with one annotation per item. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] +pub struct History { + entries: Vec>, +} + +impl History { + /// Create an empty history. + pub fn new() -> Self { + Self { + entries: Vec::new(), + } + } + + /// Build history from already annotated entries, preserving order. + pub fn from_entries(entries: Vec>) -> Self { + Self { entries } + } + + /// Replace all entries as one restore/rebuild operation and return the old entries. + pub fn replace_entries(&mut self, entries: Vec>) -> Vec> { + std::mem::replace(&mut self.entries, entries) + } + + /// Borrow annotated entries. + pub fn entries(&self) -> &[HistoryEntry] { + &self.entries + } + + /// Mutably borrow annotated entries for host-owned rebuild operations. + pub fn entries_mut(&mut self) -> &mut [HistoryEntry] { + &mut self.entries + } + + /// Consume the history into annotated entries. + pub fn into_entries(self) -> Vec> { + self.entries + } + + /// Number of entries. + pub fn len(&self) -> usize { + self.entries.len() + } + + /// Whether the history is empty. + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Iterate over annotated entries. + pub fn iter(&self) -> impl ExactSizeIterator> { + self.entries.iter() + } + + /// Iterate over provider-visible items only. + pub fn items(&self) -> impl ExactSizeIterator { + self.entries.iter().map(|entry| &entry.item) + } + + /// Clone provider-visible items into a request-local projection. + pub fn items_cloned(&self) -> Vec { + self.items().cloned().collect() + } + + /// Append an already annotated entry. + pub fn push_entry(&mut self, entry: HistoryEntry) { + self.entries.push(entry); + } + + /// Append many already annotated entries. + pub fn extend_entries(&mut self, entries: impl IntoIterator>) { + self.entries.extend(entries); + } + + /// Commit one item through a trusted annotation callback before it becomes live. + /// + /// The callback may durably persist the item and returns the annotation that + /// must be stored with it. If the callback fails, the history is left unchanged. + pub fn append_with( + &mut self, + item: Item, + annotate: &mut impl FnMut(&Item) -> Result, + ) -> Result<(), String> { + let annotation = annotate(&item)?; + self.entries.push(HistoryEntry { item, annotation }); + Ok(()) + } + + /// Commit items through a trusted annotation callback before they become live. + /// + /// Items before a failure remain appended; the failing item and later items do + /// not enter history. This mirrors append-only durable logs where each accepted + /// item is already committed before the next item is attempted. + pub fn extend_with( + &mut self, + items: impl IntoIterator, + annotate: &mut impl FnMut(&Item) -> Result, + ) -> Result<(), String> { + for item in items { + self.append_with(item, annotate)?; + } + Ok(()) + } + + /// Truncate entries, preserving item+annotation pairing for retained entries. + pub fn truncate(&mut self, len: usize) { + self.entries.truncate(len); + } + + /// Clear all entries. + pub fn clear(&mut self) { + self.entries.clear(); + } +} + +impl History<()> { + /// Build unit-annotated history from provider-visible items. + pub fn from_items(items: Vec) -> Self { + Self { + entries: items.into_iter().map(HistoryEntry::from_item).collect(), + } + } + + /// Replace history from provider-visible items using unit annotations. + pub fn replace_items(&mut self, items: Vec) -> Vec> { + self.replace_entries(items.into_iter().map(HistoryEntry::from_item).collect()) + } + + /// Append one item with unit annotation. + pub fn push(&mut self, item: Item) { + self.entries.push(HistoryEntry::from_item(item)); + } + + /// Append items with unit annotations. + pub fn extend_items(&mut self, items: impl IntoIterator) { + self.entries + .extend(items.into_iter().map(HistoryEntry::from_item)); + } +} + +impl IntoIterator for History { + type Item = HistoryEntry; + type IntoIter = std::vec::IntoIter>; + + fn into_iter(self) -> Self::IntoIter { + self.entries.into_iter() + } +} + +impl<'a, A> IntoIterator for &'a History { + type Item = &'a HistoryEntry; + type IntoIter = std::slice::Iter<'a, HistoryEntry>; + + fn into_iter(self) -> Self::IntoIter { + self.entries.iter() + } +} diff --git a/crates/agen/src/lib.rs b/crates/agen/src/lib.rs index 29769286..e6ea9953 100644 --- a/crates/agen/src/lib.rs +++ b/crates/agen/src/lib.rs @@ -2,6 +2,7 @@ mod engine; mod handler; +mod history; mod message; pub(crate) mod callback; @@ -20,13 +21,18 @@ pub mod usage_record; pub use agen_macros::{description, tool, tool_registry}; pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope}; pub use engine::{ - Engine, EngineConfig, EngineError, EngineResult, EngineRunOutput, LlmRetryNotice, - ToolRegistryError, + Engine, EngineConfig, EngineError, EngineResult, EngineRunExit, EngineRunOutput, + LlmRetryNotice, StopReason, ToolRegistryError, }; pub use handler::ToolUseBlockStart; +pub use history::{History, HistoryEntry}; pub use interceptor::Interceptor; pub use message::{ContentPart, Item, Message, Role}; -pub use tool::{ToolCall, ToolExecutionContext, ToolOutputLimits, ToolResult}; +pub use tool::{ + ToolCall, ToolExecutionContext, ToolExecutionHandle, ToolExecutionPolicy, + ToolExecutionTerminal, ToolExecutionTerminalFuture, ToolOutputLimits, ToolResult, + ToolResultDisposition, +}; pub use usage_record::UsageRecord; /// Implementation dependencies used by code generated from `agen` macros. diff --git a/crates/agen/src/llm_client/error.rs b/crates/agen/src/llm_client/error.rs index d5723646..355940fa 100644 --- a/crates/agen/src/llm_client/error.rs +++ b/crates/agen/src/llm_client/error.rs @@ -18,6 +18,9 @@ pub enum ClientError { message: String, retry_after: Option, }, + /// The provider rejected the request because it exceeded the model context window. + /// Classified only from a structured provider error code, never message text. + ContextWindowExceeded, /// A request lifecycle phase exceeded its hard timeout. Timeout { phase: &'static str, @@ -48,6 +51,7 @@ impl fmt::Display for ClientError { } write!(f, ": {}", message) } + ClientError::ContextWindowExceeded => write!(f, "Model context window reached"), ClientError::Timeout { phase, timeout } => { write!(f, "{phase} timed out after {}s", timeout.as_secs()) } @@ -112,7 +116,10 @@ pub fn is_retryable(error: &ClientError) -> bool { ClientError::Api { status: None, .. } => false, ClientError::Timeout { .. } => true, ClientError::Http(e) => e.is_connect() || e.is_timeout(), - ClientError::Json(_) | ClientError::Sse(_) | ClientError::Config(_) => false, + ClientError::ContextWindowExceeded + | ClientError::Json(_) + | ClientError::Sse(_) + | ClientError::Config(_) => false, } } diff --git a/crates/agen/src/llm_client/transport.rs b/crates/agen/src/llm_client/transport.rs index f71edfc0..1198a9e2 100644 --- a/crates/agen/src/llm_client/transport.rs +++ b/crates/agen/src/llm_client/transport.rs @@ -431,13 +431,7 @@ fn api_error_code(error: &ClientError) -> Option<&str> { } fn is_context_length_exceeded(error: &ClientError) -> bool { - match error { - ClientError::Api { code, message, .. } => { - code.as_deref() == Some("context_length_exceeded") - || message.contains("context_length_exceeded") - } - _ => false, - } + matches!(error, ClientError::ContextWindowExceeded) } async fn response_with_timeout( @@ -487,6 +481,9 @@ async fn classify_error_response(resp: reqwest::Response) -> ClientError { .and_then(|v| v.as_str()) .unwrap_or(&text) .to_string(); + if code.as_deref() == Some("context_length_exceeded") { + return ClientError::ContextWindowExceeded; + } ClientError::Api { status: Some(status), code, diff --git a/crates/agen/src/llm_client/types.rs b/crates/agen/src/llm_client/types.rs index dd2d3dc9..101e9773 100644 --- a/crates/agen/src/llm_client/types.rs +++ b/crates/agen/src/llm_client/types.rs @@ -9,7 +9,7 @@ use std::{fmt, sync::Arc}; -use crate::tool::Attachment; +use crate::tool::{Attachment, ToolResultDisposition}; use base64::Engine as _; use serde::{Deserialize, Serialize}; @@ -121,6 +121,9 @@ pub enum Item { /// Detailed output (removed by pruning when old enough) #[serde(default, skip_serializing_if = "Option::is_none")] content: Option, + /// Typed terminal state used for replay and recovery. + #[serde(default, skip_serializing_if = "ToolResultDisposition::is_success")] + disposition: ToolResultDisposition, /// Whether the tool result represents an execution error. #[serde(default, skip_serializing_if = "is_false")] is_error: bool, @@ -261,7 +264,17 @@ impl Item { content: Option, is_error: bool, ) -> Self { - Self::tool_result_item_with_attachments(call_id, summary, content, is_error, Vec::new()) + Self::tool_result_item_with_disposition_and_attachments( + call_id, + summary, + content, + if is_error { + ToolResultDisposition::Error + } else { + ToolResultDisposition::Success + }, + Vec::new(), + ) } /// Create a tool result item with durable, prunable structured attachments. @@ -272,11 +285,33 @@ impl Item { is_error: bool, attachments: Vec, ) -> Self { + Self::tool_result_item_with_disposition_and_attachments( + call_id, + summary, + content, + if is_error { + ToolResultDisposition::Error + } else { + ToolResultDisposition::Success + }, + attachments, + ) + } + + pub fn tool_result_item_with_disposition_and_attachments( + call_id: impl Into, + summary: impl Into, + content: Option, + disposition: ToolResultDisposition, + attachments: Vec, + ) -> Self { + let is_error = !disposition.is_success(); Self::ToolResult { id: None, call_id: call_id.into(), summary: summary.into(), content, + disposition, is_error, attachments, } diff --git a/crates/agen/src/state.rs b/crates/agen/src/state.rs index 5d1b5172..3ca3c14a 100644 --- a/crates/agen/src/state.rs +++ b/crates/agen/src/state.rs @@ -19,7 +19,7 @@ mod private { /// - Editing message history (add, delete, clear) /// - Registering tools and hooks /// -/// Can transition to [`Locked`] state via `Engine::lock()`. +/// Can transition to [`Locked`] state via `Engine::lock(&history)`. /// /// # Examples /// diff --git a/crates/agen/src/tool.rs b/crates/agen/src/tool.rs index 6fe66fa6..7557d978 100644 --- a/crates/agen/src/tool.rs +++ b/crates/agen/src/tool.rs @@ -3,7 +3,14 @@ //! Traits for defining tools callable by LLM. //! Usually auto-implemented using the `#[tool]` macro. -use std::{collections::HashMap, fmt, sync::Arc}; +use std::{ + collections::HashMap, + fmt, + future::Future, + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; use async_trait::async_trait; use base64::{Engine as _, engine::general_purpose::STANDARD}; @@ -23,6 +30,12 @@ pub enum ToolError { /// Internal error #[error("Internal error: {0}")] Internal(String), + /// Cooperative cancellation completed with bounded terminal output. + #[error("Tool execution cancelled")] + Cancelled(ToolOutput), + /// Execution was interrupted with a confirmed bounded terminal output. + #[error("Tool execution interrupted")] + Interrupted(ToolOutput), } // ============================================================================= @@ -158,6 +171,28 @@ pub enum Attachment { Image(ImageAttachment), } +/// Terminal disposition of one started tool call. +/// +/// `Cancelled` means the tool confirmed cancellation. `OutcomeUnknown` means +/// execution stopped without confirmation, so neither completion nor side +/// effects may be inferred. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "snake_case")] +pub enum ToolResultDisposition { + #[default] + Success, + Error, + Interrupted, + Cancelled, + OutcomeUnknown, +} + +impl ToolResultDisposition { + pub const fn is_success(&self) -> bool { + matches!(self, Self::Success) + } +} + /// Tool execution result. /// /// Every output has a mandatory `summary` (1-2 lines) that persists in @@ -322,6 +357,12 @@ impl ToolExecutionContext { } } + /// Identifies one live execution attempt without making the batch id a durable + /// replay or idempotency authority. + pub fn execution_id(&self) -> String { + format!("{}:{}", self.batch_id, self.call_id) + } + /// Context for direct, non-engine calls in unit tests and low-level callers. pub fn direct() -> Self { Self::new("direct", "direct", 0) @@ -334,6 +375,142 @@ impl Default for ToolExecutionContext { } } +/// The provider-confirmed terminal result of one started tool execution. +/// +/// `OutcomeUnknown` is reserved for an execution task that had to be force-closed +/// or failed before the provider could confirm its terminal result. +#[derive(Debug)] +pub enum ToolExecutionTerminal { + Confirmed(Result), + OutcomeUnknown, +} + +/// The completion future paired with a [`ToolExecutionHandle`]. Dropping this +/// future does not drop the provider execution: the spawned execution remains +/// owned by its handle until it completes or is explicitly force-closed. +pub struct ToolExecutionTerminalFuture { + task: tokio::task::JoinHandle>, +} + +impl Future for ToolExecutionTerminalFuture { + type Output = ToolExecutionTerminal; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match Pin::new(&mut self.task).poll(cx) { + Poll::Ready(Ok(result)) => Poll::Ready(ToolExecutionTerminal::Confirmed(result)), + Poll::Ready(Err(_)) => Poll::Ready(ToolExecutionTerminal::OutcomeUnknown), + Poll::Pending => Poll::Pending, + } + } +} + +/// Live ownership and control for one started tool execution. +/// +/// Execution, cancellation, and terminal confirmation remain provider-owned: +/// this handle starts `Tool::execute`, delegates cooperative cancellation to +/// `Tool::cancel_execution`, and treats execution-future completion as the +/// provider's terminal confirmation. Agen may force-close only after its caller's +/// deadline expires, at which point the outcome is necessarily unknown. +#[derive(Clone)] +pub struct ToolExecutionHandle { + inner: Arc, +} + +struct ToolExecutionHandleInner { + tool: Arc, + context: ToolExecutionContext, + abort: tokio::task::AbortHandle, +} + +impl Drop for ToolExecutionHandleInner { + fn drop(&mut self) { + // Losing the final live owner is an explicit forced close, never a + // best-effort detached provider future. + self.abort.abort(); + } +} + +impl fmt::Debug for ToolExecutionHandle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ToolExecutionHandle") + .field("call_id", &self.inner.context.call_id) + .field("batch_id", &self.inner.context.batch_id) + .finish_non_exhaustive() + } +} + +impl ToolExecutionHandle { + pub fn start( + tool: Arc, + input_json: String, + context: ToolExecutionContext, + ) -> (Self, ToolExecutionTerminalFuture) { + let execution_tool = Arc::clone(&tool); + let execution_context = context.clone(); + let task = + tokio::spawn( + async move { execution_tool.execute(&input_json, execution_context).await }, + ); + let abort = task.abort_handle(); + ( + Self { + inner: Arc::new(ToolExecutionHandleInner { + tool, + context, + abort, + }), + }, + ToolExecutionTerminalFuture { task }, + ) + } + + pub fn context(&self) -> &ToolExecutionContext { + &self.inner.context + } + + pub async fn cancel_before(&self, deadline: tokio::time::Instant) -> Result<(), ToolError> { + match tokio::time::timeout_at( + deadline, + self.inner.tool.cancel_execution(&self.inner.context), + ) + .await + { + Ok(result) => result, + Err(_) => Err(ToolError::Internal(format!( + "tool cancellation request exceeded its deadline for call {}", + self.inner.context.call_id + ))), + } + } + + pub fn force_close(&self) { + self.inner.abort.abort(); + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ToolExecutionPolicy { + /// Time a pause waits for already-started providers to reach a natural safe + /// boundary before escalating to explicit cooperative cancellation. + pub pause_safe_boundary_timeout: std::time::Duration, + /// Maximum time allowed for a provider to accept one cooperative + /// cancellation request. + pub cancellation_request_timeout: std::time::Duration, + /// Maximum time allowed for all providers to confirm terminal results after + /// cancellation has been requested. + pub terminal_confirmation_timeout: std::time::Duration, +} + +impl Default for ToolExecutionPolicy { + fn default() -> Self { + Self { + pause_safe_boundary_timeout: std::time::Duration::from_millis(100), + cancellation_request_timeout: std::time::Duration::from_millis(100), + terminal_confirmation_timeout: std::time::Duration::from_millis(500), + } + } +} + // ============================================================================= // Tool trait // ============================================================================= @@ -402,6 +579,26 @@ pub trait Tool: Send + Sync { input_json: &str, ctx: ToolExecutionContext, ) -> Result; + + /// Request cooperative cancellation for one started call. + /// + /// Implementations that own cancellable provider operations should signal + /// every live execution identified by `call_id`, then let `execute` return + /// the confirmed bounded terminal output. Direct callers may use this + /// compatibility surface; Agen uses [`Tool::cancel_execution`] so providers + /// can bind cancellation to one exact live attempt. + async fn cancel(&self, _call_id: &str) -> Result<(), ToolError> { + Ok(()) + } + + /// Request cooperative cancellation for one exact started execution. + /// + /// The default preserves existing tools by delegating to `cancel(call_id)`. + /// Providers with their own execution registry should override this method + /// and key cancellation by [`ToolExecutionContext::execution_id`]. + async fn cancel_execution(&self, ctx: &ToolExecutionContext) -> Result<(), ToolError> { + self.cancel(&ctx.call_id).await + } } // ============================================================================= @@ -429,6 +626,9 @@ pub struct ToolCall { pub struct ToolResult { /// Corresponding tool call ID pub tool_use_id: String, + /// Typed terminal state. + #[serde(default, skip_serializing_if = "ToolResultDisposition::is_success")] + pub disposition: ToolResultDisposition, /// Short summary (always kept in history) pub summary: String, /// Detailed output (prunable) @@ -445,11 +645,20 @@ pub struct ToolResult { impl ToolResult { /// Create a success result from a [`ToolOutput`]. pub fn from_output(tool_use_id: impl Into, output: ToolOutput) -> Self { + Self::from_output_with_disposition(tool_use_id, output, ToolResultDisposition::Success) + } + + pub fn from_output_with_disposition( + tool_use_id: impl Into, + output: ToolOutput, + disposition: ToolResultDisposition, + ) -> Self { Self { tool_use_id: tool_use_id.into(), + disposition, summary: output.summary, content: output.content, - is_error: false, + is_error: !disposition.is_success(), attachments: output.attachments, } } @@ -458,12 +667,28 @@ impl ToolResult { pub fn error(tool_use_id: impl Into, message: impl Into) -> Self { Self { tool_use_id: tool_use_id.into(), + disposition: ToolResultDisposition::Error, summary: message.into(), content: None, is_error: true, attachments: Vec::new(), } } + + /// Close an execution whose completion and side effects cannot be confirmed. + pub fn outcome_unknown(tool_use_id: impl Into) -> Self { + Self { + tool_use_id: tool_use_id.into(), + disposition: ToolResultDisposition::OutcomeUnknown, + summary: "Tool execution outcome unknown".to_string(), + content: Some( + "Execution was interrupted before completion could be confirmed. Completion and side effects are unknown." + .to_string(), + ), + is_error: true, + attachments: Vec::new(), + } + } } #[cfg(test)] diff --git a/crates/agen/tests/annotated_history_test.rs b/crates/agen/tests/annotated_history_test.rs new file mode 100644 index 00000000..09a5019d --- /dev/null +++ b/crates/agen/tests/annotated_history_test.rs @@ -0,0 +1,84 @@ +mod common; + +use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; +use agen::{Engine, EngineError, History, HistoryEntry, Item, Role}; +use common::MockLlmClient; + +fn completed_text_events(text: &str) -> Vec { + vec![ + Event::text_block_start(0), + Event::text_delta(0, text), + Event::text_block_stop(0, None), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] +} + +#[tokio::test] +async fn run_preserves_item_annotations_without_projecting_them() { + let client = MockLlmClient::new(completed_text_events("assistant reply")); + let engine = Engine::<_, agen::state::Mutable, String>::new_annotated(client); + let mut history = History::::new(); + let mut next = 0usize; + let mut annotate = |item: &Item| { + next += 1; + let kind = match item { + Item::Message { role, .. } => match role { + Role::User => "user", + Role::Assistant => "assistant", + Role::System => "system", + }, + Item::ToolCall { .. } => "tool_call", + Item::ToolResult { .. } => "tool_result", + Item::Reasoning { .. } => "reasoning", + }; + Ok(format!("{next}:{kind}")) + }; + + let output = engine + .run_with_annotation(&mut history, "hello", &mut annotate) + .await; + + assert!(matches!(output.result, agen::EngineRunExit::Finished)); + assert_eq!(history.len(), 2); + assert_eq!(history.entries()[0].annotation, "1:user"); + assert_eq!(history.entries()[1].annotation, "2:assistant"); + assert_eq!(history.items_cloned().len(), 2); +} + +#[test] +fn append_failure_does_not_make_item_live() { + let client = MockLlmClient::new(vec![]); + let mut engine = Engine::<_, agen::state::Mutable, usize>::new_annotated(client); + let mut history = History::::new(); + let mut fail = |_item: &Item| Err("commit failed".to_string()); + + let err = engine + .append_history_with(&mut history, [Item::user_message("uncommitted")], &mut fail) + .unwrap_err(); + + assert!(matches!(err, EngineError::HistoryAppend(message) if message == "commit failed")); + assert!(history.is_empty()); +} + +#[test] +fn replacement_keeps_items_and_annotations_together() { + let mut history = History::from_entries(vec![ + HistoryEntry::new(Item::user_message("old"), "old-ann".to_string()), + HistoryEntry::new(Item::user_message("second"), "second-ann".to_string()), + ]); + + history.truncate(1); + assert_eq!(history.entries()[0].item.as_text(), Some("old")); + assert_eq!(history.entries()[0].annotation, "old-ann"); + + let previous = history.replace_entries(vec![HistoryEntry::new( + Item::user_message("restored"), + "restored-ann".to_string(), + )]); + + assert_eq!(previous.len(), 1); + assert_eq!(history.entries()[0].item.as_text(), Some("restored")); + assert_eq!(history.entries()[0].annotation, "restored-ann"); +} diff --git a/crates/agen/tests/callback_test.rs b/crates/agen/tests/callback_test.rs index d3f9b60e..1c3535f8 100644 --- a/crates/agen/tests/callback_test.rs +++ b/crates/agen/tests/callback_test.rs @@ -58,6 +58,7 @@ async fn test_callback_llm_retry_event() { max_attempts: 2, total_timeout: Duration::from_secs(1), }); + let mut history = agen::History::new(); let notices = Arc::new(Mutex::new(Vec::new())); let sink = notices.clone(); @@ -65,8 +66,11 @@ async fn test_callback_llm_retry_event() { sink.lock().unwrap().push((llm_call, notice.clone())); }); - let result = engine.run("retry once").await; - assert!(result.is_ok(), "engine should succeed after one retry"); + let result = engine.run(&mut history, "retry once").await; + assert!( + matches!(result.result, agen::EngineRunExit::Finished), + "engine should succeed after one retry" + ); let notices = notices.lock().unwrap(); assert_eq!(notices.len(), 1); @@ -91,6 +95,7 @@ async fn test_callback_text_block_events() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history = agen::History::new(); let text_deltas = Arc::new(Mutex::new(Vec::new())); let text_completes = Arc::new(Mutex::new(Vec::new())); @@ -108,9 +113,12 @@ async fn test_callback_text_block_events() { }); }); - // Mutable::run consumes self, returns (Locked, EngineResult) - let result = engine.run("Greet me").await; - assert!(result.is_ok(), "Engine should complete"); + // Mutable::run consumes self, returns (Locked, EngineRunExit) + let result = engine.run(&mut history, "Greet me").await; + assert!( + matches!(result.result, agen::EngineRunExit::Finished), + "Engine should complete" + ); let deltas = text_deltas.lock().unwrap(); assert_eq!(deltas.len(), 2); @@ -137,6 +145,7 @@ async fn test_callback_tool_call_complete() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history = agen::History::new(); let tool_starts = Arc::new(Mutex::new(Vec::<(String, String)>::new())); let tool_completes = Arc::new(Mutex::new(Vec::new())); @@ -154,8 +163,8 @@ async fn test_callback_tool_call_complete() { }); }); - // Mutable::run consumes self, returns (Locked, EngineResult) - let _ = engine.run("Weather please").await; + // Mutable::run consumes self, returns (Locked, EngineRunExit) + let _ = engine.run(&mut history, "Weather please").await; let starts = tool_starts.lock().unwrap(); assert_eq!(starts.len(), 1); @@ -183,6 +192,7 @@ async fn test_callback_turn_events() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history = agen::History::new(); let turn_starts = Arc::new(Mutex::new(Vec::new())); let turn_ends = Arc::new(Mutex::new(Vec::new())); @@ -197,9 +207,9 @@ async fn test_callback_turn_events() { ends.lock().unwrap().push(turn); }); - // Mutable::run consumes self, returns (Locked, EngineResult) - let result = engine.run("Do something").await; - assert!(result.is_ok()); + // Mutable::run consumes self, returns (Locked, EngineRunExit) + let result = engine.run(&mut history, "Do something").await; + assert!(matches!(result.result, agen::EngineRunExit::Finished)); let starts = turn_starts.lock().unwrap(); let ends = turn_ends.lock().unwrap(); @@ -254,6 +264,7 @@ async fn test_callback_tool_result_events() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history = agen::History::new(); engine.register_tool(fixed_tool( "fixed", @@ -276,7 +287,7 @@ async fn test_callback_tool_result_events() { )); }); - let _ = engine.run("call it").await; + let _ = engine.run(&mut history, "call it").await; let observed = captured.lock().unwrap(); assert_eq!(observed.len(), 1); @@ -330,6 +341,7 @@ async fn test_callback_tool_result_error_path() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history = agen::History::new(); engine.register_tool(erroring_tool("erroring", "boom")); @@ -345,7 +357,7 @@ async fn test_callback_tool_result_error_path() { )); }); - let _ = engine.run("fail it").await; + let _ = engine.run(&mut history, "fail it").await; let observed = captured.lock().unwrap(); assert_eq!(observed.len(), 1); @@ -374,6 +386,7 @@ async fn test_callback_usage_events() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history = agen::History::new(); let usage_events = Arc::new(Mutex::new(Vec::new())); @@ -382,8 +395,8 @@ async fn test_callback_usage_events() { usages.lock().unwrap().push(event.clone()); }); - // Mutable::run consumes self, returns (Locked, EngineResult) - let _ = engine.run("Hello").await; + // Mutable::run consumes self, returns (Locked, EngineRunExit) + let _ = engine.run(&mut history, "Hello").await; let usages = usage_events.lock().unwrap(); assert_eq!(usages.len(), 1); diff --git a/crates/agen/tests/common/mod.rs b/crates/agen/tests/common/mod.rs index 81902e78..5f4d4e94 100644 --- a/crates/agen/tests/common/mod.rs +++ b/crates/agen/tests/common/mod.rs @@ -19,6 +19,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; pub struct MockLlmClient { responses: Arc>>, call_count: Arc, + requests: Arc>>, } impl MockLlmClient { @@ -30,6 +31,7 @@ impl MockLlmClient { Self { responses: Arc::new(responses), call_count: Arc::new(AtomicUsize::new(0)), + requests: Arc::new(Mutex::new(Vec::new())), } } @@ -41,6 +43,10 @@ impl MockLlmClient { pub fn event_count(&self) -> usize { self.responses.iter().map(|v| v.len()).sum() } + + pub fn requests(&self) -> Vec { + self.requests.lock().unwrap().clone() + } } #[async_trait] @@ -51,8 +57,9 @@ impl LlmClient for MockLlmClient { async fn stream( &self, - _request: Request, + request: Request, ) -> Result> + Send>>, ClientError> { + self.requests.lock().unwrap().push(request); let count = self.call_count.fetch_add(1, Ordering::SeqCst); if count >= self.responses.len() { return Err(ClientError::Api { diff --git a/crates/agen/tests/engine_fixtures.rs b/crates/agen/tests/engine_fixtures.rs index 1ab9e9ef..e3f63e97 100644 --- a/crates/agen/tests/engine_fixtures.rs +++ b/crates/agen/tests/engine_fixtures.rs @@ -134,11 +134,15 @@ async fn test_engine_simple_text_response() { let client = MockLlmClient::from_fixture(&fixture_path).unwrap(); let engine = Engine::new(client); + let mut history = agen::History::new(); // Send a simple message (Mutable::run consumes self, returns tuple) - let result = engine.run("Hello").await; + let result = engine.run(&mut history, "Hello").await; - assert!(result.is_ok(), "Engine should complete successfully"); + assert!( + matches!(result.result, agen::EngineRunExit::Finished), + "Engine should complete successfully" + ); } /// Verify that Engine can correctly process responses containing tool calls @@ -156,6 +160,7 @@ async fn test_engine_tool_call() { let client = MockLlmClient::from_fixture(&fixture_path).unwrap(); let mut engine = Engine::new(client); + let mut history = agen::History::new(); // Register tool let weather_tool = MockWeatherTool::new(); @@ -163,7 +168,9 @@ async fn test_engine_tool_call() { engine.register_tool(weather_tool.definition()); // Send message (Mutable::run consumes self, returns tuple) - let _result = engine.run("What's the weather in Tokyo?").await; + let _result = engine + .run(&mut history, "What's the weather in Tokyo?") + .await; // Verify tool was called // Note: max_turns=1 so no request is sent after tool result @@ -195,11 +202,15 @@ async fn test_engine_with_programmatic_events() { let client = MockLlmClient::new(events); let engine = Engine::new(client); + let mut history = agen::History::new(); // Mutable::run consumes self, returns tuple - let result = engine.run("Greet me").await; + let result = engine.run(&mut history, "Greet me").await; - assert!(result.is_ok(), "Engine should complete successfully"); + assert!( + matches!(result.result, agen::EngineRunExit::Finished), + "Engine should complete successfully" + ); } /// Verify that ToolCallCollector correctly collects ToolCall from ToolUse block events diff --git a/crates/agen/tests/engine_state_test.rs b/crates/agen/tests/engine_state_test.rs index 5bffd7e4..c1bfe966 100644 --- a/crates/agen/tests/engine_state_test.rs +++ b/crates/agen/tests/engine_state_test.rs @@ -14,7 +14,7 @@ use agen::interceptor::{ }; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; -use agen::{Engine, EngineError, EngineResult}; +use agen::{Engine, EngineError, EngineRunExit, History, StopReason}; use async_trait::async_trait; use common::MockLlmClient; @@ -42,36 +42,37 @@ fn test_mutable_set_system_prompt() { fn test_mutable_history_manipulation() { let client = MockLlmClient::new(vec![]); let mut engine = Engine::new(client); + let mut history: History = History::new(); // Initial state is empty - assert!(engine.history().is_empty()); + assert!(history.is_empty()); // Add to history engine - .append_history(vec![Item::user_message("Hello")]) + .append_history(&mut history, vec![Item::user_message("Hello")]) .unwrap(); engine - .append_history(vec![Item::assistant_message("Hi there!")]) + .append_history(&mut history, vec![Item::assistant_message("Hi there!")]) .unwrap(); - assert_eq!(engine.history().len(), 2); + assert_eq!(history.len(), 2); // Append to history via the callback-aware API. engine - .append_history(vec![Item::user_message("How are you?")]) + .append_history(&mut history, vec![Item::user_message("How are you?")]) .unwrap(); - assert_eq!(engine.history().len(), 3); + assert_eq!(history.len(), 3); // Clear history - engine.clear_history(); - assert!(engine.history().is_empty()); + engine.clear_history(&mut history); + assert!(history.is_empty()); // Set history let items = vec![ Item::user_message("Test"), Item::assistant_message("Response"), ]; - engine.set_history(items); - assert_eq!(engine.history().len(), 2); + engine.set_history(&mut history, items); + assert_eq!(history.len(), 2); } /// Verify that Engine can be constructed using builder pattern @@ -79,9 +80,10 @@ fn test_mutable_history_manipulation() { fn test_mutable_builder_pattern() { let client = MockLlmClient::new(vec![]); let engine = Engine::new(client).system_prompt("System prompt"); + let history: History = History::new(); assert_eq!(engine.get_system_prompt(), Some("System prompt")); - assert!(engine.history().is_empty()); + assert!(history.is_empty()); } /// Verify that multiple items can be added with append_history and callbacks fire. @@ -91,6 +93,7 @@ fn test_mutable_append_history() { let observed = Arc::new(Mutex::new(Vec::new())); let observed_for_callback = Arc::clone(&observed); let mut engine = Engine::new(client); + let mut history: History = History::new(); engine.on_history_append(move |item| { if let Some(text) = item.as_text() { observed_for_callback.lock().unwrap().push(text.to_string()); @@ -99,18 +102,21 @@ fn test_mutable_append_history() { }); engine - .append_history(vec![Item::user_message("First")]) + .append_history(&mut history, vec![Item::user_message("First")]) .unwrap(); engine - .append_history(vec![ - Item::assistant_message("Response 1"), - Item::user_message("Second"), - Item::assistant_message("Response 2"), - ]) + .append_history( + &mut history, + vec![ + Item::assistant_message("Response 1"), + Item::user_message("Second"), + Item::assistant_message("Response 2"), + ], + ) .unwrap(); - assert_eq!(engine.history().len(), 4); + assert_eq!(history.len(), 4); assert_eq!( observed.lock().unwrap().as_slice(), ["First", "Response 1", "Second", "Response 2"] @@ -185,6 +191,7 @@ async fn history_append_failure_stops_before_tool_execution() { ]); let tool = CountingTool::new("count_tool"); let mut engine = Engine::new(client); + let mut history: History = History::new(); engine.register_tool(tool.definition()); engine.on_history_append(|item| { if item.is_tool_call() { @@ -194,15 +201,15 @@ async fn history_append_failure_stops_before_tool_execution() { } }); - let mut engine = engine.lock(); - let error = engine.run("use the tool").await.unwrap_err(); + let mut engine = engine.lock(&history); + let exit = engine.run(&mut history, "use the tool").await; assert!( - matches!(error, EngineError::HistoryAppend(ref message) if message == "simulated ENOSPC") + matches!(exit, EngineRunExit::Interrupted(StopReason::Unexpected(EngineError::HistoryAppend(ref message))) if message == "simulated ENOSPC") ); assert_eq!(tool.call_count(), 0); - assert_eq!(engine.history().len(), 1); - assert_eq!(engine.history()[0].as_text(), Some("use the tool")); + assert_eq!(history.len(), 1); + assert_eq!(history.entries()[0].item.as_text(), Some("use the tool")); } // ============================================================================= @@ -214,21 +221,22 @@ async fn history_append_failure_stops_before_tool_execution() { fn test_lock_transition() { let client = MockLlmClient::new(vec![]); let mut engine = Engine::new(client); + let mut history: History = History::new(); engine.set_system_prompt("System"); engine - .append_history(vec![Item::user_message("Hello")]) + .append_history(&mut history, vec![Item::user_message("Hello")]) .unwrap(); engine - .append_history(vec![Item::assistant_message("Hi")]) + .append_history(&mut history, vec![Item::assistant_message("Hi")]) .unwrap(); // Lock - let locked_engine = engine.lock(); + let locked_engine = engine.lock(&history); // History and system prompt are still accessible in Locked state assert_eq!(locked_engine.get_system_prompt(), Some("System")); - assert_eq!(locked_engine.history().len(), 2); + assert_eq!(history.len(), 2); assert_eq!(locked_engine.locked_prefix_len(), 2); } @@ -237,21 +245,22 @@ fn test_lock_transition() { fn test_unlock_transition() { let client = MockLlmClient::new(vec![]); let mut engine = Engine::new(client); + let mut history: History = History::new(); engine - .append_history(vec![Item::user_message("Hello")]) + .append_history(&mut history, vec![Item::user_message("Hello")]) .unwrap(); - let locked_engine = engine.lock(); + let locked_engine = engine.lock(&history); // Unlock let mut engine = locked_engine.unlock(); // History operations are available again in Mutable state engine - .append_history(vec![Item::assistant_message("Hi")]) + .append_history(&mut history, vec![Item::assistant_message("Hi")]) .unwrap(); - engine.clear_history(); - assert!(engine.history().is_empty()); + engine.clear_history(&mut history); + assert!(history.is_empty()); } // ============================================================================= @@ -272,20 +281,20 @@ async fn test_mutable_run_updates_history() -> Result<(), EngineError> { let client = MockLlmClient::new(events); let engine = Engine::new(client); + let mut history: History = History::new(); // Execute (Mutable::run consumes self, returns EngineRunOutput) - let out = engine.run("Hi there").await?; - let engine = out.engine; + let _out = engine.run(&mut history, "Hi there").await; // History is updated - let history = engine.history(); + let entries = history.entries(); assert_eq!(history.len(), 2); // user + assistant // User message - assert_eq!(history[0].as_text(), Some("Hi there")); + assert_eq!(entries[0].item.as_text(), Some("Hi there")); // Assistant message - assert_eq!(history[1].as_text(), Some("Hello, I'm an assistant!")); + assert_eq!(entries[1].item.as_text(), Some("Hello, I'm an assistant!")); Ok(()) } @@ -316,35 +325,36 @@ async fn test_locked_multi_turn_history_accumulation() { ]); let engine = Engine::new(client).system_prompt("You are helpful."); + let mut history: History = History::new(); // Lock (after setting system prompt) - let mut locked_engine = engine.lock(); + let mut locked_engine = engine.lock(&history); assert_eq!(locked_engine.locked_prefix_len(), 0); // No items yet // Turn 1 - let result1 = locked_engine.run("Hello!").await; - assert!(result1.is_ok()); - assert_eq!(locked_engine.history().len(), 2); // user + assistant + let result1 = locked_engine.run(&mut history, "Hello!").await; + assert!(matches!(result1, EngineRunExit::Finished)); + assert_eq!(history.len(), 2); // user + assistant // Turn 2 - let result2 = locked_engine.run("Can you help me?").await; - assert!(result2.is_ok()); - assert_eq!(locked_engine.history().len(), 4); // 2 * (user + assistant) + let result2 = locked_engine.run(&mut history, "Can you help me?").await; + assert!(matches!(result2, EngineRunExit::Finished)); + assert_eq!(history.len(), 4); // 2 * (user + assistant) // Verify history contents - let history = locked_engine.history(); + let entries = history.entries(); // Turn 1 user message - assert_eq!(history[0].as_text(), Some("Hello!")); + assert_eq!(entries[0].item.as_text(), Some("Hello!")); // Turn 1 assistant message - assert_eq!(history[1].as_text(), Some("Nice to meet you!")); + assert_eq!(entries[1].item.as_text(), Some("Nice to meet you!")); // Turn 2 user message - assert_eq!(history[2].as_text(), Some("Can you help me?")); + assert_eq!(entries[2].item.as_text(), Some("Can you help me?")); // Turn 2 assistant message - assert_eq!(history[3].as_text(), Some("I can help with that.")); + assert_eq!(entries[3].item.as_text(), Some("I can help with that.")); } /// Verify that locked_prefix_len correctly records history length at lock time @@ -370,26 +380,33 @@ async fn test_locked_prefix_len_tracking() { ]); let mut engine = Engine::new(client); + let mut history: History = History::new(); // Add items beforehand engine - .append_history(vec![Item::user_message("Pre-existing message 1")]) + .append_history( + &mut history, + vec![Item::user_message("Pre-existing message 1")], + ) .unwrap(); engine - .append_history(vec![Item::assistant_message("Pre-existing response 1")]) + .append_history( + &mut history, + vec![Item::assistant_message("Pre-existing response 1")], + ) .unwrap(); - assert_eq!(engine.history().len(), 2); + assert_eq!(history.len(), 2); // Lock - let mut locked_engine = engine.lock(); + let mut locked_engine = engine.lock(&history); assert_eq!(locked_engine.locked_prefix_len(), 2); // 2 items at lock time // Execute turn - locked_engine.run("New message").await.unwrap(); + locked_engine.run(&mut history, "New message").await; // History grows but locked_prefix_len remains unchanged - assert_eq!(locked_engine.history().len(), 4); // 2 + 2 + assert_eq!(history.len(), 4); // 2 + 2 assert_eq!(locked_engine.locked_prefix_len(), 2); // Unchanged } @@ -416,18 +433,22 @@ async fn test_turn_count_increment() -> Result<(), EngineError> { ]); let engine = Engine::new(client); + let mut history: History = History::new(); assert_eq!(engine.turn_count(), 0); assert_eq!(engine.llm_call_count(), 0); // First run consumes Mutable, returns EngineRunOutput - let mut engine = engine.run("First").await?.engine; + let mut engine = engine.run(&mut history, "First").await.engine; assert_eq!(engine.turn_count(), 1); // Retry not yet implemented → AgentTurn:LlmCall is 1:1. assert_eq!(engine.llm_call_count(), 1); // Subsequent runs on Locked take &mut self - engine.run("Second").await?; + assert!(matches!( + engine.run(&mut history, "Second").await, + EngineRunExit::Finished + )); assert_eq!(engine.turn_count(), 2); assert_eq!(engine.llm_call_count(), 2); @@ -447,28 +468,29 @@ async fn test_unlock_edit_relock() { ]]); let mut engine = Engine::new(client); + let mut history: History = History::new(); engine - .append_history(vec![ - Item::user_message("Hello"), - Item::assistant_message("Hi"), - ]) + .append_history( + &mut history, + vec![Item::user_message("Hello"), Item::assistant_message("Hi")], + ) .unwrap(); // Lock -> Unlock - let locked = engine.lock(); + let locked = engine.lock(&history); assert_eq!(locked.locked_prefix_len(), 2); let mut unlocked = locked.unlock(); // Edit history - unlocked.clear_history(); + unlocked.clear_history(&mut history); unlocked - .append_history(vec![Item::user_message("Fresh start")]) + .append_history(&mut history, vec![Item::user_message("Fresh start")]) .unwrap(); // Re-lock - let relocked = unlocked.lock(); - assert_eq!(relocked.history().len(), 1); + let relocked = unlocked.lock(&history); + assert_eq!(history.len(), 1); assert_eq!(relocked.locked_prefix_len(), 1); } @@ -511,19 +533,26 @@ async fn test_lock_unlock_relock_tools_remain_effective() { ]); let mut engine = Engine::new(client); + let mut history: History = History::new(); let tool_a = CountingTool::new("tool_a"); engine.register_tool(tool_a.definition()); - let mut locked = engine.lock(); - locked.run("first").await.expect("first run"); + let mut locked = engine.lock(&history); + assert!(matches!( + locked.run(&mut history, "first").await, + EngineRunExit::Finished + )); assert_eq!(tool_a.call_count(), 1, "tool_a should be called once"); let mut unlocked = locked.unlock(); let tool_b = CountingTool::new("tool_b"); unlocked.register_tool(tool_b.definition()); - let mut relocked = unlocked.lock(); - relocked.run("second").await.expect("second run"); + let mut relocked = unlocked.lock(&history); + assert!(matches!( + relocked.run(&mut history, "second").await, + EngineRunExit::Finished + )); assert_eq!(tool_a.call_count(), 1, "tool_a should not be called again"); assert_eq!(tool_b.call_count(), 1, "tool_b should be called once"); @@ -538,8 +567,9 @@ async fn test_lock_unlock_relock_tools_remain_effective() { fn test_system_prompt_preserved_in_locked_state() { let client = MockLlmClient::new(vec![]); let engine = Engine::new(client).system_prompt("Important system prompt"); + let history: History = History::new(); - let locked = engine.lock(); + let locked = engine.lock(&history); assert_eq!(locked.get_system_prompt(), Some("Important system prompt")); let unlocked = locked.unlock(); @@ -554,14 +584,15 @@ fn test_system_prompt_preserved_in_locked_state() { fn test_system_prompt_change_after_unlock() { let client = MockLlmClient::new(vec![]); let engine = Engine::new(client).system_prompt("Original prompt"); + let history: History = History::new(); - let locked = engine.lock(); + let locked = engine.lock(&history); let mut unlocked = locked.unlock(); unlocked.set_system_prompt("New prompt"); assert_eq!(unlocked.get_system_prompt(), Some("New prompt")); - let relocked = unlocked.lock(); + let relocked = unlocked.lock(&history); assert_eq!(relocked.get_system_prompt(), Some("New prompt")); } @@ -623,40 +654,55 @@ impl Interceptor for ContinueTurnOnce { #[tokio::test] async fn max_turns_is_scoped_to_each_fresh_run() { + let mut history: History = History::new(); let responses = vec![completed_text_events(), completed_text_events()]; let mut engine = Engine::new(MockLlmClient::with_responses(responses)); engine.set_max_turns(Some(1)); - let mut engine = engine.lock(); + let mut engine = engine.lock(&history); - assert_eq!(engine.run("first").await.unwrap(), EngineResult::Finished); + assert!(matches!( + engine.run(&mut history, "first").await, + EngineRunExit::Finished + )); assert_eq!(engine.turn_count(), 1); assert_eq!(engine.active_run_turn_count(), None); - assert_eq!(engine.run("second").await.unwrap(), EngineResult::Finished); + assert!(matches!( + engine.run(&mut history, "second").await, + EngineRunExit::Finished + )); assert_eq!(engine.turn_count(), 2); assert_eq!(engine.active_run_turn_count(), None); } #[tokio::test] async fn yielded_resume_keeps_the_same_unspent_turn_budget() { + let mut history: History = History::new(); let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); engine.set_max_turns(Some(1)); engine.set_interceptor(YieldOnce { calls: AtomicUsize::new(0), }); - let mut engine = engine.lock(); + let mut engine = engine.lock(&history); - assert_eq!(engine.run("start").await.unwrap(), EngineResult::Yielded); + assert!(matches!( + engine.run(&mut history, "start").await, + EngineRunExit::Yielded + )); assert_eq!(engine.turn_count(), 0); assert_eq!(engine.active_run_turn_count(), Some(0)); - assert_eq!(engine.resume().await.unwrap(), EngineResult::Finished); + assert!(matches!( + engine.resume(&mut history).await, + EngineRunExit::Finished + )); assert_eq!(engine.turn_count(), 1); assert_eq!(engine.active_run_turn_count(), None); } #[tokio::test] async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() { + let mut history: History = History::new(); let events = vec![ Event::tool_use_start(0, "call_1", "count_tool"), Event::tool_input_delta(0, "{}"), @@ -672,14 +718,20 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() { engine.set_interceptor(PauseToolOnce { calls: AtomicUsize::new(0), }); - let mut engine = engine.lock(); + let mut engine = engine.lock(&history); - assert_eq!(engine.run("call it").await.unwrap(), EngineResult::Paused); + assert!(matches!( + engine.run(&mut history, "call it").await, + EngineRunExit::Paused + )); assert_eq!(engine.turn_count(), 1); assert_eq!(engine.active_run_turn_count(), Some(1)); assert_eq!(tool.call_count(), 0); - assert_eq!(engine.resume().await.unwrap(), EngineResult::LimitReached); + assert!(matches!( + engine.resume(&mut history).await, + EngineRunExit::Interrupted(StopReason::LimitReached) + )); assert_eq!(engine.turn_count(), 1); assert_eq!(engine.active_run_turn_count(), None); assert_eq!(tool.call_count(), 1, "the consumed turn's tool still runs"); @@ -687,6 +739,7 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() { #[tokio::test] async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() { + let mut history: History = History::new(); let tool_events = vec![ Event::tool_use_start(0, "call_1", "count_tool"), Event::tool_input_delta(0, "{}"), @@ -703,12 +756,18 @@ async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() { engine.set_interceptor(PauseToolOnce { calls: AtomicUsize::new(0), }); - let mut engine = engine.lock(); + let mut engine = engine.lock(&history); - assert_eq!(engine.run("pause").await.unwrap(), EngineResult::Paused); + assert!(matches!( + engine.run(&mut history, "pause").await, + EngineRunExit::Paused + )); assert_eq!(engine.active_run_turn_count(), Some(1)); - assert_eq!(engine.run("replace").await.unwrap(), EngineResult::Finished); + assert!(matches!( + engine.run(&mut history, "replace").await, + EngineRunExit::Finished + )); assert_eq!(engine.turn_count(), 2); assert_eq!(engine.active_run_turn_count(), None); assert_eq!(tool.call_count(), 1, "pending-tool semantics are unchanged"); @@ -716,17 +775,18 @@ async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() { #[tokio::test] async fn interceptor_continuation_consumes_the_logical_run_budget() { + let mut history: History = History::new(); let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); engine.set_max_turns(Some(1)); engine.set_interceptor(ContinueTurnOnce { calls: AtomicUsize::new(0), }); - let mut engine = engine.lock(); + let mut engine = engine.lock(&history); - assert_eq!( - engine.run("start").await.unwrap(), - EngineResult::LimitReached - ); + assert!(matches!( + engine.run(&mut history, "start").await, + EngineRunExit::Interrupted(StopReason::LimitReached) + )); assert_eq!(engine.turn_count(), 1); assert_eq!(engine.llm_call_count(), 1); assert_eq!(engine.active_run_turn_count(), None); @@ -734,14 +794,17 @@ async fn interceptor_continuation_consumes_the_logical_run_budget() { #[tokio::test] async fn restored_active_run_budget_is_enforced_before_another_llm_call() { + let mut history: History = History::new(); let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); engine.set_max_turns(Some(1)); engine.set_turn_count(7); - engine.set_last_run_interrupted(true); engine.set_active_run_turn_count(Some(1)); - let mut engine = engine.lock(); + let mut engine = engine.lock(&history); - assert_eq!(engine.resume().await.unwrap(), EngineResult::LimitReached); + assert!(matches!( + engine.resume(&mut history).await, + EngineRunExit::Interrupted(StopReason::LimitReached) + )); assert_eq!(engine.turn_count(), 7); assert_eq!(engine.llm_call_count(), 0); assert_eq!(engine.active_run_turn_count(), None); diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index a7067377..61f29600 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -6,12 +6,13 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; -use agen::Engine; use agen::interceptor::{Interceptor, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::tool::{ Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult, + ToolResultDisposition, }; +use agen::{Engine, History, Item, ToolExecutionPolicy}; use async_trait::async_trait; mod common; @@ -70,6 +71,144 @@ impl Tool for SlowTool { } } +#[derive(Clone)] +struct FirstAttemptHangsTool { + calls: Arc, +} + +impl FirstAttemptHangsTool { + fn new() -> Self { + Self { + calls: Arc::new(AtomicUsize::new(0)), + } + } + + fn definition(&self) -> ToolDefinition { + let tool = self.clone(); + Arc::new(move || { + let meta = ToolMeta::new("hang_once") + .description("Hangs on the first execution attempt") + .input_schema(serde_json::json!({"type": "object"})); + (meta, Arc::new(tool.clone()) as Arc) + }) + } + + fn call_count(&self) -> usize { + self.calls.load(Ordering::SeqCst) + } +} + +#[async_trait] +impl Tool for FirstAttemptHangsTool { + async fn execute( + &self, + _input_json: &str, + _ctx: ToolExecutionContext, + ) -> Result { + let attempt = self.calls.fetch_add(1, Ordering::SeqCst); + if attempt == 0 { + std::future::pending::<()>().await; + } + Ok("completed on retry".to_string().into()) + } +} + +#[derive(Clone)] +struct CooperativeCancelTool { + calls: Arc, + cancelled: Arc, +} + +impl CooperativeCancelTool { + fn new() -> Self { + Self { + calls: Arc::new(AtomicUsize::new(0)), + cancelled: Arc::new(tokio::sync::Notify::new()), + } + } + + fn definition(&self) -> ToolDefinition { + let tool = self.clone(); + Arc::new(move || { + let meta = ToolMeta::new("cooperative") + .description("Returns bounded progress after cancellation") + .input_schema(serde_json::json!({"type": "object"})); + (meta, Arc::new(tool.clone()) as Arc) + }) + } +} + +#[async_trait] +impl Tool for CooperativeCancelTool { + async fn execute( + &self, + _input_json: &str, + _ctx: ToolExecutionContext, + ) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + self.cancelled.notified().await; + Err(ToolError::Cancelled(ToolOutput { + summary: "cooperative command cancelled".to_string(), + content: Some("stdout before cancellation\nstderr before cancellation".to_string()), + attachments: Vec::new(), + })) + } + + async fn cancel(&self, _call_id: &str) -> Result<(), ToolError> { + self.cancelled.notify_one(); + Ok(()) + } +} + +#[derive(Clone)] +struct SafePauseTool { + calls: Arc, + cancellations: Arc, + release: Arc, +} + +impl SafePauseTool { + fn new() -> Self { + Self { + calls: Arc::new(AtomicUsize::new(0)), + cancellations: Arc::new(AtomicUsize::new(0)), + release: Arc::new(tokio::sync::Notify::new()), + } + } + + fn definition(&self) -> ToolDefinition { + let tool = self.clone(); + Arc::new(move || { + let meta = ToolMeta::new("safe_pause") + .description("Waits for a safe-boundary release") + .input_schema(serde_json::json!({"type": "object"})); + (meta, Arc::new(tool.clone()) as Arc) + }) + } +} + +#[async_trait] +impl Tool for SafePauseTool { + async fn execute( + &self, + _input_json: &str, + _ctx: ToolExecutionContext, + ) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + self.release.notified().await; + Ok(ToolOutput { + summary: "safe-boundary complete".to_string(), + content: Some("safe-boundary complete".to_string()), + attachments: Vec::new(), + }) + } + + async fn cancel(&self, _call_id: &str) -> Result<(), ToolError> { + self.cancellations.fetch_add(1, Ordering::SeqCst); + Ok(()) + } +} + #[derive(Clone)] struct ContextRecordingTool { name: String, @@ -145,6 +284,7 @@ async fn test_parallel_tool_execution() { ], ]); let mut engine = Engine::new(client); + let mut history: History = History::new(); let tool1 = SlowTool::new("slow_tool_1", 100); let tool2 = SlowTool::new("slow_tool_2", 100); let tool3 = SlowTool::new("slow_tool_3", 100); @@ -159,7 +299,7 @@ async fn test_parallel_tool_execution() { let start = Instant::now(); // Mutable::run consumes self, returns (Locked, EngineResult) - let _result = engine.run("Run all tools").await; + let _result = engine.run(&mut history, "Run all tools").await; let elapsed = start.elapsed(); // Verify all tools were called @@ -178,6 +318,450 @@ async fn test_parallel_tool_execution() { println!("Parallel execution completed in {:?}", elapsed); } +#[tokio::test] +async fn completed_results_commit_before_publish_without_waiting_for_siblings() { + let client = MockLlmClient::with_responses(vec![ + vec![ + Event::tool_use_start(0, "call_slow", "slow_first"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::tool_use_start(1, "call_fast", "fast_second"), + Event::tool_input_delta(1, r#"{}"#), + Event::tool_use_stop(1), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ], + vec![ + Event::text_block_start(0), + Event::text_delta(0, "Done"), + Event::text_block_stop(0, None), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ], + ]); + let client_probe = client.clone(); + let mut engine = Engine::new(client); + engine.register_tool(SlowTool::new("slow_first", 100).definition()); + engine.register_tool(SlowTool::new("fast_second", 5).definition()); + + let observed = Arc::new(Mutex::new(Vec::::new())); + let published = observed.clone(); + engine.on_tool_result(move |result| { + published + .lock() + .unwrap() + .push(format!("publish:{}", result.tool_use_id)); + }); + + let committed = observed.clone(); + let mut annotate = move |item: &Item| { + if let Item::ToolResult { call_id, .. } = item { + committed.lock().unwrap().push(format!("commit:{call_id}")); + } + Ok(()) + }; + let mut history = History::new(); + let _ = engine + .run_with_annotation(&mut history, "run both", &mut annotate) + .await; + observed.lock().unwrap().push("run-returned".to_string()); + + assert_eq!( + observed.lock().unwrap().as_slice(), + [ + "commit:call_fast", + "publish:call_fast", + "commit:call_slow", + "publish:call_slow", + "run-returned", + ] + ); + + let committed_order: Vec<_> = history + .iter() + .filter_map(|entry| match &entry.item { + Item::ToolResult { call_id, .. } => Some(call_id.as_str()), + _ => None, + }) + .collect(); + assert_eq!(committed_order, ["call_fast", "call_slow"]); + + let requests = client_probe.requests(); + let projected_order: Vec<_> = requests[1] + .items + .iter() + .filter_map(|item| match item { + Item::ToolResult { call_id, .. } => Some(call_id.as_str()), + _ => None, + }) + .collect(); + assert_eq!(projected_order, ["call_slow", "call_fast"]); +} + +#[tokio::test] +async fn cancellation_preserves_completed_results_and_resume_skips_them() { + let client = MockLlmClient::with_responses(vec![ + vec![ + Event::tool_use_start(0, "call_hang", "hang_once"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::tool_use_start(1, "call_fast_a", "fast_a"), + Event::tool_input_delta(1, r#"{}"#), + Event::tool_use_stop(1), + Event::tool_use_start(2, "call_fast_b", "fast_b"), + Event::tool_input_delta(2, r#"{}"#), + Event::tool_use_stop(2), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ], + vec![ + Event::text_block_start(0), + Event::text_delta(0, "Recovered"), + Event::text_block_stop(0, None), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ], + ]); + let mut engine = Engine::new(client); + let hanging = FirstAttemptHangsTool::new(); + let fast_a = SlowTool::new("fast_a", 1); + let fast_b = SlowTool::new("fast_b", 2); + engine.register_tool(hanging.definition()); + engine.register_tool(fast_a.definition()); + engine.register_tool(fast_b.definition()); + + let cancel = engine.cancel_sender(); + let cancel_task = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(30)).await; + cancel.send(()).await.unwrap(); + }); + let mut history = History::new(); + let output = engine.run(&mut history, "start").await; + let mut engine = output.engine; + cancel_task.await.unwrap(); + + let completed_before_resume = history + .iter() + .filter(|entry| { + matches!( + &entry.item, + Item::ToolResult { call_id, .. } + if call_id == "call_fast_a" || call_id == "call_fast_b" + ) + }) + .count(); + let unknown_before_resume = history + .iter() + .filter(|entry| { + matches!( + &entry.item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call_hang" + ) + }) + .count(); + assert_eq!(completed_before_resume, 2); + assert_eq!(unknown_before_resume, 1); + assert_eq!(fast_a.call_count(), 1); + assert_eq!(fast_b.call_count(), 1); + assert_eq!(hanging.call_count(), 1); + + let _ = engine.resume(&mut history).await; + + assert_eq!( + fast_a.call_count(), + 1, + "completed call must not be re-executed" + ); + assert_eq!( + fast_b.call_count(), + 1, + "completed call must not be re-executed" + ); + assert_eq!( + hanging.call_count(), + 1, + "OutcomeUnknown is terminal and must not be re-executed" + ); + let completed_after_resume = history + .iter() + .filter(|entry| { + matches!( + &entry.item, + Item::ToolResult { call_id, .. } + if call_id == "call_fast_a" || call_id == "call_fast_b" + ) + }) + .count(); + assert_eq!(completed_after_resume, 2); + assert_eq!( + history + .iter() + .filter(|entry| { + matches!( + &entry.item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call_hang" + ) + }) + .count(), + 1 + ); +} + +#[tokio::test] +async fn cooperative_cancellation_commits_bounded_terminal_output() { + let client = MockLlmClient::with_responses(vec![vec![ + Event::tool_use_start(0, "call_cooperative", "cooperative"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]]); + let mut engine = Engine::new(client); + let tool = CooperativeCancelTool::new(); + engine.register_tool(tool.definition()); + let observed = Arc::new(Mutex::new(Vec::<&'static str>::new())); + let published = observed.clone(); + engine.on_tool_result(move |_| published.lock().unwrap().push("published")); + let committed = observed.clone(); + let mut annotate = move |item: &Item| { + if matches!(item, Item::ToolResult { .. }) { + committed.lock().unwrap().push("committed"); + } + Ok(()) + }; + + let cancel = engine.cancel_sender(); + let cancel_task = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(30)).await; + cancel.send(()).await.unwrap(); + }); + let mut history = History::new(); + let output = engine + .run_with_annotation(&mut history, "start", &mut annotate) + .await; + observed.lock().unwrap().push("run-returned"); + cancel_task.await.unwrap(); + + assert_eq!( + observed.lock().unwrap().as_slice(), + ["committed", "published", "run-returned"] + ); + assert_eq!(tool.calls.load(Ordering::SeqCst), 1); + let terminal: Vec<_> = history + .iter() + .filter_map(|entry| match &entry.item { + Item::ToolResult { + call_id, + disposition, + content, + .. + } if call_id == "call_cooperative" => Some((*disposition, content.as_deref())), + _ => None, + }) + .collect(); + assert_eq!(terminal.len(), 1); + assert_eq!(terminal[0].0, ToolResultDisposition::Cancelled); + assert_eq!( + terminal[0].1, + Some("stdout before cancellation\nstderr before cancellation") + ); + assert!(matches!( + output.result, + agen::EngineRunExit::Interrupted(agen::StopReason::Cancelled) + )); +} + +#[tokio::test] +async fn pause_waits_for_started_tool_terminal_without_cancelling_provider() { + let client = MockLlmClient::with_responses(vec![vec![ + Event::tool_use_start(0, "call_safe_pause", "safe_pause"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]]); + let mut engine = Engine::new(client); + let tool = SafePauseTool::new(); + engine.register_tool(tool.definition()); + + let pause = engine.pause_sender(); + let calls = Arc::clone(&tool.calls); + let release = Arc::clone(&tool.release); + let control = tokio::spawn(async move { + tokio::time::timeout(Duration::from_secs(1), async { + while calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("tool execution starts"); + pause.send(()).await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + release.notify_one(); + }); + + let started_at = std::time::Instant::now(); + let mut history = History::new(); + let output = engine.run(&mut history, "pause safely").await; + control.await.unwrap(); + + assert!(started_at.elapsed() >= Duration::from_millis(50)); + assert_eq!(tool.calls.load(Ordering::SeqCst), 1); + assert_eq!(tool.cancellations.load(Ordering::SeqCst), 0); + assert!(matches!(output.result, agen::EngineRunExit::Paused)); + assert!(history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::Success, + .. + } if call_id == "call_safe_pause" + ))); +} + +#[tokio::test] +async fn pause_escalates_to_explicit_cancel_and_confirm_after_safe_boundary_deadline() { + let client = MockLlmClient::with_responses(vec![vec![ + Event::tool_use_start(0, "call_pause_cancel", "cooperative"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]]); + let mut engine = Engine::new(client); + engine.set_tool_execution_policy(ToolExecutionPolicy { + pause_safe_boundary_timeout: Duration::from_millis(20), + cancellation_request_timeout: Duration::from_millis(50), + terminal_confirmation_timeout: Duration::from_millis(100), + }); + let tool = CooperativeCancelTool::new(); + engine.register_tool(tool.definition()); + + let pause = engine.pause_sender(); + let calls = Arc::clone(&tool.calls); + let control = tokio::spawn(async move { + tokio::time::timeout(Duration::from_secs(1), async { + while calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("tool execution starts"); + pause.send(()).await.unwrap(); + }); + + let mut history = History::new(); + let output = engine.run(&mut history, "pause with escalation").await; + control.await.unwrap(); + + assert!(matches!(output.result, agen::EngineRunExit::Paused)); + assert!(history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::Cancelled, + .. + } if call_id == "call_pause_cancel" + ))); +} + +#[tokio::test] +async fn cancellation_completion_race_commits_one_terminal_output() { + for iteration in 0..24u64 { + let client = MockLlmClient::with_responses(vec![ + vec![ + Event::tool_use_start(0, "call_racy", "racy"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ], + vec![Event::Status(StatusEvent { + status: ResponseStatus::Completed, + })], + ]); + let mut engine = Engine::new(client); + let delay = 2 + iteration % 3; + let tool = SlowTool::new("racy", delay); + engine.register_tool(tool.definition()); + let cancel = engine.cancel_sender(); + let cancel_task = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(delay)).await; + let _ = cancel.send(()).await; + }); + + let mut history = History::new(); + let _ = engine.run(&mut history, "race").await; + cancel_task.await.unwrap(); + let terminal_count = history + .iter() + .filter(|entry| { + matches!( + &entry.item, + Item::ToolResult { call_id, .. } if call_id == "call_racy" + ) + }) + .count(); + assert_eq!(terminal_count, 1, "iteration {iteration}"); + assert_eq!(tool.call_count(), 1, "iteration {iteration}"); + } +} + +#[tokio::test] +async fn tool_result_commit_failure_prevents_publication() { + let client = MockLlmClient::with_responses(vec![vec![ + Event::tool_use_start(0, "call_fast", "fast"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]]); + let mut engine = Engine::new(client); + engine.register_tool(SlowTool::new("fast", 1).definition()); + + let published = Arc::new(AtomicUsize::new(0)); + let published_probe = published.clone(); + engine.on_tool_result(move |_| { + published_probe.fetch_add(1, Ordering::SeqCst); + }); + + let mut history = History::new(); + let mut reject_tool_result = |item: &Item| { + if matches!(item, Item::ToolResult { .. }) { + Err("session log unavailable".to_string()) + } else { + Ok(()) + } + }; + let _ = engine + .run_with_annotation(&mut history, "start", &mut reject_tool_result) + .await; + + assert_eq!(published.load(Ordering::SeqCst), 0); + assert!( + history + .iter() + .all(|entry| !matches!(entry.item, Item::ToolResult { .. })) + ); +} + #[tokio::test] async fn test_tool_execution_context_order_and_batch_id() { let client = MockLlmClient::with_responses(vec![ @@ -205,13 +789,14 @@ async fn test_tool_execution_context_order_and_batch_id() { ], ]); let mut engine = Engine::new(client); + let mut history: History = History::new(); let contexts = Arc::new(Mutex::new(Vec::new())); engine.register_tool(ContextRecordingTool::new("record_a", contexts.clone()).definition()); engine.register_tool(ContextRecordingTool::new("record_b", contexts.clone()).definition()); engine.register_tool(ContextRecordingTool::new("record_c", contexts.clone()).definition()); - let _ = engine.run("record contexts").await; + let _ = engine.run(&mut history, "record contexts").await; let mut contexts = contexts.lock().unwrap().clone(); contexts.sort_by_key(|ctx| ctx.call_index); @@ -256,11 +841,12 @@ async fn test_tool_execution_context_batch_id_changes_between_batches() { ], ]); let mut engine = Engine::new(client); + let mut history: History = History::new(); let contexts = Arc::new(Mutex::new(Vec::new())); engine.register_tool(ContextRecordingTool::new("record", contexts.clone()).definition()); - let _ = engine.run("record batches").await; + let _ = engine.run(&mut history, "record batches").await; let contexts = contexts.lock().unwrap().clone(); assert_eq!(contexts.len(), 2); @@ -298,6 +884,7 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() { ], ]); let mut engine = Engine::new(client); + let mut history: History = History::new(); let executed_contexts = Arc::new(Mutex::new(Vec::new())); let pre_contexts = Arc::new(Mutex::new(Vec::new())); let post_contexts = Arc::new(Mutex::new(Vec::new())); @@ -344,7 +931,9 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() { post_contexts: post_contexts.clone(), }); - let _ = engine.run("record skipped and synthetic contexts").await; + let _ = engine + .run(&mut history, "record skipped and synthetic contexts") + .await; let mut pre_contexts = pre_contexts.lock().unwrap().clone(); pre_contexts.sort_by_key(|ctx| ctx.call_index); @@ -389,6 +978,7 @@ async fn test_before_tool_call_skip() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history: History = History::new(); let allowed_tool = SlowTool::new("allowed_tool", 10); let blocked_tool = SlowTool::new("blocked_tool", 10); @@ -416,7 +1006,7 @@ async fn test_before_tool_call_skip() { engine.set_interceptor(BlockingPolicy); // Mutable::run consumes self, returns (Locked, EngineResult) - let _result = engine.run("Test hook").await; + let _result = engine.run(&mut history, "Test hook").await; // allowed_tool is called, but blocked_tool is not assert_eq!( @@ -457,6 +1047,7 @@ async fn test_post_tool_call_modification() { ]); let mut engine = Engine::new(client); + let mut history: History = History::new(); #[derive(Clone)] struct SimpleTool; @@ -503,9 +1094,12 @@ async fn test_post_tool_call_modification() { }); // Mutable::run consumes self, returns (Locked, EngineResult) - let result = engine.run("Test modification").await; + let result = engine.run(&mut history, "Test modification").await; - assert!(result.is_ok(), "Engine should complete"); + assert!( + matches!(result.result, agen::EngineRunExit::Finished), + "Engine should complete" + ); // Verify hook was called and content was modified let content = modified_content.lock().unwrap().clone(); @@ -540,6 +1134,7 @@ async fn test_before_tool_call_synthetic_result_committed() { ], ]); let mut engine = Engine::new(client); + let mut history: History = History::new(); let blocked_tool = SlowTool::new("blocked_tool", 10); let blocked_clone = blocked_tool.clone(); engine.register_tool(blocked_tool.definition()); @@ -558,10 +1153,10 @@ async fn test_before_tool_call_synthetic_result_committed() { engine.set_interceptor(SyntheticPolicy); - let result = engine.run("Test synthetic result").await.unwrap(); + let _result = engine.run(&mut history, "Test synthetic result").await; assert_eq!(blocked_clone.call_count(), 0, "Blocked tool should not run"); - assert!(result.engine.history().iter().any(|item| matches!( + assert!(history.items().any(|item| matches!( item, agen::Item::ToolResult { call_id, @@ -571,3 +1166,76 @@ async fn test_before_tool_call_synthetic_result_committed() { } if call_id == "call_1" && summary == "permission denied" ))); } + +#[tokio::test] +async fn post_tool_abort_commits_confirmed_result_before_stopping_run() { + let client = MockLlmClient::new(vec![ + Event::tool_use_start(0, "call_confirmed", "confirmed"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]); + let mut engine = Engine::new(client); + let tool = SlowTool::new("confirmed", 1); + engine.register_tool(tool.definition()); + + struct AbortAfterResult; + #[async_trait] + impl Interceptor for AbortAfterResult { + async fn post_tool_call(&self, _info: &mut ToolResultInfo) -> PostToolAction { + PostToolAction::Abort("policy stopped the run".to_string()) + } + } + engine.set_interceptor(AbortAfterResult); + + let observed = Arc::new(Mutex::new(Vec::<&'static str>::new())); + let published = observed.clone(); + engine.on_tool_result(move |_| published.lock().unwrap().push("published")); + let committed = observed.clone(); + let mut annotate = move |item: &Item| { + if matches!(item, Item::ToolResult { .. }) { + committed.lock().unwrap().push("committed"); + } + Ok(()) + }; + + let mut history = History::new(); + let output = engine + .run_with_annotation(&mut history, "run confirmed tool", &mut annotate) + .await; + observed.lock().unwrap().push("run-returned"); + + assert_eq!(tool.call_count(), 1); + assert_eq!( + observed.lock().unwrap().as_slice(), + ["committed", "published", "run-returned"] + ); + assert!(matches!( + output.result, + agen::EngineRunExit::Interrupted(agen::StopReason::Unexpected( + agen::EngineError::Aborted(ref reason) + )) if reason == "policy stopped the run" + )); + let terminal: Vec<_> = history + .iter() + .filter_map(|entry| match &entry.item { + Item::ToolResult { + call_id, + disposition, + .. + } if call_id == "call_confirmed" => Some(*disposition), + _ => None, + }) + .collect(); + assert_eq!(terminal, [ToolResultDisposition::Success]); + assert!(!history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call_confirmed" + ))); +} diff --git a/crates/agen/tests/reasoning_round_trip_test.rs b/crates/agen/tests/reasoning_round_trip_test.rs index 3b37a258..21e83f2a 100644 --- a/crates/agen/tests/reasoning_round_trip_test.rs +++ b/crates/agen/tests/reasoning_round_trip_test.rs @@ -13,12 +13,12 @@ mod common; -use agen::Engine; use agen::Item; use agen::llm_client::event::{ BlockMetadata, BlockStart, BlockStop, BlockType, Event, ReasoningBlockData, ResponseStatus, StatusEvent, }; +use agen::{Engine, History}; use common::MockLlmClient; fn reasoning_block(text: impl Into, data: ReasoningBlockData) -> Vec { @@ -65,15 +65,15 @@ async fn anthropic_thinking_round_trips_signature_into_history() { ]); let client = MockLlmClient::new(events); let engine = Engine::new(client); - let out = engine.run("question?").await.expect("run ok"); - let engine = out.engine; + let mut history: History = History::new(); + let _out = engine.run(&mut history, "question?").await; - let history = engine.history(); + let entries = history.entries(); // user / reasoning / assistant_message assert_eq!(history.len(), 3, "history: {history:?}"); - assert!(matches!(history[0], Item::Message { .. })); - match &history[1] { + assert!(matches!(entries[0].item, Item::Message { .. })); + match &entries[1].item { Item::Reasoning { text, signature, .. } => { @@ -82,7 +82,7 @@ async fn anthropic_thinking_round_trips_signature_into_history() { } other => panic!("expected Reasoning, got {other:?}"), } - assert_eq!(history[2].as_text(), Some("Here's the answer")); + assert_eq!(entries[2].item.as_text(), Some("Here's the answer")); } /// OpenAI Responses 風: encrypted_content + summary を持った reasoning が @@ -109,11 +109,11 @@ async fn openai_reasoning_round_trips_encrypted_and_summary() { ]); let client = MockLlmClient::new(events); let engine = Engine::new(client); - let out = engine.run("q").await.expect("run ok"); - let engine = out.engine; + let mut history: History = History::new(); + let _out = engine.run(&mut history, "q").await; - let history = engine.history(); - match &history[1] { + let entries = history.entries(); + match &entries[1].item { Item::Reasoning { text, summary, @@ -155,13 +155,13 @@ async fn reasoning_precedes_text_in_assistant_burst() { })); let client = MockLlmClient::new(events); let engine = Engine::new(client); - let out = engine.run("q").await.expect("run ok"); - let engine = out.engine; + let mut history: History = History::new(); + let _out = engine.run(&mut history, "q").await; - let history = engine.history(); + let entries = history.entries(); // user / reasoning(先頭) / assistant_message - assert!(matches!(history[1], Item::Reasoning { .. })); - assert_eq!(history[2].as_text(), Some("intermediate")); + assert!(matches!(entries[1].item, Item::Reasoning { .. })); + assert_eq!(entries[2].item.as_text(), Some("intermediate")); } /// resume シナリオ: history.json 由来の Item::Reasoning(signature) を Engine に @@ -207,14 +207,18 @@ async fn injected_reasoning_survives_into_outgoing_request() { }; let mut engine = Engine::new(client); + let mut history: History = History::new(); // resume: 既存 history を流し込む - engine.set_history(vec![ - Item::user_message("prior question"), - Item::reasoning("prior thinking").with_signature("SIG-PRIOR"), - Item::assistant_message("prior answer"), - ]); + engine.set_history( + &mut history, + vec![ + Item::user_message("prior question"), + Item::reasoning("prior thinking").with_signature("SIG-PRIOR"), + Item::assistant_message("prior answer"), + ], + ); - let _ = engine.run("follow up").await.expect("run ok"); + let _ = engine.run(&mut history, "follow up").await; let req = captured .lock() diff --git a/crates/agen/tests/ui/locked_register_tool.rs b/crates/agen/tests/ui/locked_register_tool.rs index e30c93b7..e48b282e 100644 --- a/crates/agen/tests/ui/locked_register_tool.rs +++ b/crates/agen/tests/ui/locked_register_tool.rs @@ -1,4 +1,4 @@ -use agen::Engine; +use agen::{Engine, History}; use agen::llm_client::capability::{ CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport, }; @@ -22,7 +22,8 @@ fn main() { cap, ); let engine = Engine::new(client); - let mut locked = engine.lock(); + let history = History::new(); + let mut locked = engine.lock(&history); let def: agen::tool::ToolDefinition = Arc::new(|| panic!("unused")); let _ = locked.register_tool(def); } diff --git a/crates/agen/tests/ui/locked_register_tool.stderr b/crates/agen/tests/ui/locked_register_tool.stderr index 394c8cc1..c0edc4ff 100644 --- a/crates/agen/tests/ui/locked_register_tool.stderr +++ b/crates/agen/tests/ui/locked_register_tool.stderr @@ -1,8 +1,8 @@ error[E0599]: no method named `register_tool` found for struct `Engine, Locked>` in the current scope - --> tests/ui/locked_register_tool.rs:27:20 + --> tests/ui/locked_register_tool.rs:28:20 | -27 | let _ = locked.register_tool(def); +28 | let _ = locked.register_tool(def); | ^^^^^^^^^^^^^ method not found in `Engine, Locked>` | = note: the method was found for - - `Engine` + - `Engine` diff --git a/crates/fs-operation/src/lib.rs b/crates/fs-operation/src/lib.rs index f63a1ce3..53999081 100644 --- a/crates/fs-operation/src/lib.rs +++ b/crates/fs-operation/src/lib.rs @@ -279,4 +279,58 @@ mod tests { assert_eq!(grep.matched_files, 2); assert!(!grep.output.contains("c.txt")); } + + #[test] + fn grep_content_groups_lines_by_file_and_marks_matches() { + let temp = tempfile::tempdir().unwrap(); + std::fs::write( + temp.path().join("first.txt"), + "before\nneedle one\nafter\nomitted one\nomitted two\nbefore distant\nneedle distant\nafter distant\n", + ) + .unwrap(); + std::fs::write(temp.path().join("second.txt"), "needle two\n").unwrap(); + let root = temp.path().canonicalize().unwrap(); + let readable = RootAccess(root.clone()); + + let grep = run_grep( + &root, + root.clone(), + GrepRequest { + pattern: "needle".to_string(), + path: FsPath::root(), + glob: Some("*.txt".to_string()), + output_mode: GrepOutputMode::Content, + case_insensitive: false, + before_context: 1, + after_context: 1, + multiline: false, + file_type: None, + limit: 20, + offset: 0, + }, + &readable, + ) + .unwrap(); + + assert_eq!(grep.match_count, 3); + assert_eq!(grep.matched_files, 2); + assert_eq!( + grep.output, + concat!( + "first.txt\n", + " 1 │ before\n", + " > 2 │ needle one\n", + " 3 │ after\n", + " …\n", + " 6 │ before distant\n", + " > 7 │ needle distant\n", + " 8 │ after distant\n", + "\n", + "second.txt\n", + " > 1 │ needle two\n", + ) + ); + assert_eq!(grep.output.matches("first.txt").count(), 1); + assert_eq!(grep.output.matches("second.txt").count(), 1); + } } diff --git a/crates/fs-operation/src/search.rs b/crates/fs-operation/src/search.rs index 74a66d27..54420f6f 100644 --- a/crates/fs-operation/src/search.rs +++ b/crates/fs-operation/src/search.rs @@ -1,3 +1,5 @@ +use std::collections::BTreeMap; +use std::fmt::Write as _; use std::path::{Path, PathBuf}; use crate::FsAccessPolicy; @@ -57,20 +59,11 @@ impl GrepReport { } } GrepOutputMode::Content => { - for line in &self.lines { - let separator = if line.is_match { ':' } else { '-' }; - let path = logical_display(root, &line.path); - if self.show_line_numbers - && let Some(number) = line.line_number - { - output.push_str(&format!( - "{path}{separator}{number}{separator}{}\n", - line.text - )); - } else { - output.push_str(&format!("{path}{separator}{}\n", line.text)); - } - } + output.push_str(&render_content_lines( + root, + &self.lines, + self.show_line_numbers, + )); } } GrepResult { @@ -82,6 +75,48 @@ impl GrepReport { } } +fn render_content_lines(root: &Path, lines: &[ContentLine], show_line_numbers: bool) -> String { + let mut grouped = BTreeMap::<&Path, Vec<&ContentLine>>::new(); + for line in lines { + grouped.entry(&line.path).or_default().push(line); + } + + let mut output = String::new(); + for (file_index, (path, file_lines)) in grouped.into_iter().enumerate() { + if file_index > 0 { + output.push('\n'); + } + let _ = writeln!(output, "{}", logical_display(root, path)); + + let number_width = file_lines + .iter() + .filter_map(|line| line.line_number) + .map(|number| number.to_string().len()) + .max() + .unwrap_or(1); + let mut previous_line_end = None; + for line in file_lines { + if let (Some(previous_end), Some(number)) = (previous_line_end, line.line_number) + && number > previous_end + { + let _ = writeln!(output, " …"); + } + + let marker = if line.is_match { '>' } else { ' ' }; + if show_line_numbers && let Some(number) = line.line_number { + let _ = writeln!(output, " {marker} {number:>number_width$} │ {}", line.text); + } else { + let _ = writeln!(output, " {marker} │ {}", line.text); + } + previous_line_end = line + .line_number + .map(|number| number + line.text.split('\n').count() as u64); + } + } + + output +} + fn logical_display(root: &Path, path: &Path) -> String { path.strip_prefix(root) .unwrap_or(path) diff --git a/crates/memory/src/extract/payload.rs b/crates/memory/src/extract/payload.rs index e4780c9e..42f87a29 100644 --- a/crates/memory/src/extract/payload.rs +++ b/crates/memory/src/extract/payload.rs @@ -9,7 +9,7 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use crate::schema::{EvidenceKind, SourceEvidenceRef, SourceRef}; +use crate::schema::{EvidenceKind, EvidenceOrigin, SourceEvidenceRef, SourceRef}; /// Current flat staging schema version. pub const STAGING_SCHEMA_VERSION: u32 = 2; @@ -80,6 +80,8 @@ pub struct StagingEvidence { #[serde(default, skip_serializing_if = "Option::is_none")] pub entry_range: Option<[u64; 2]>, #[serde(default, skip_serializing_if = "Option::is_none")] + pub origin: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub excerpt: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub summary: Option, @@ -159,6 +161,7 @@ mod tests { id: "E001".into(), kind: EvidenceKind::new(EvidenceKind::MESSAGE), entry_range: Some([10, 12]), + origin: None, excerpt: Some("extract candidate taxonomy".into()), summary: Some("User and assistant discussed staging kinds".into()), }; diff --git a/crates/memory/src/schema/common.rs b/crates/memory/src/schema/common.rs index 8a5f4085..a5396632 100644 --- a/crates/memory/src/schema/common.rs +++ b/crates/memory/src/schema/common.rs @@ -67,6 +67,40 @@ impl EvidenceKind { } } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum EvidenceOriginKind { + HumanInput, + WorkerInput, + FlowInstruction, + BackendInstruction, + ModelOutput, + ToolOutput, + DerivedSummary, + LegacyUnknown, +} + +/// Bounded origin snapshot attached to extraction evidence. This is audit +/// metadata only and cannot authorize Workspace operations. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +pub struct EvidenceOrigin { + pub kind: EvidenceOriginKind, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub account_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub workspace_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub runtime_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub worker_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub flow_selector: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub flow_definition_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub flow_definition_revision: Option, +} + /// Host-resolved source/evidence metadata for an individual staging claim. /// /// This deliberately stores only bounded anchor metadata: stable ids, entry @@ -86,6 +120,9 @@ pub struct SourceEvidenceRef { /// Host-assigned evidence id within the referenced evidence set. #[serde(default, skip_serializing_if = "Option::is_none")] pub evidence_id: Option, + /// Trusted typed origin snapshot for this logical evidence entry. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub origin: Option, /// Extensible evidence kind tag. #[serde(default, skip_serializing_if = "Option::is_none")] pub evidence_kind: Option, diff --git a/crates/memory/src/schema/mod.rs b/crates/memory/src/schema/mod.rs index 67155c5d..b2e99833 100644 --- a/crates/memory/src/schema/mod.rs +++ b/crates/memory/src/schema/mod.rs @@ -10,7 +10,10 @@ mod decision; mod request; mod summary; -pub use common::{EvidenceKind, Frontmatter, SourceEvidenceRef, SourceRef, split_frontmatter}; +pub use common::{ + EvidenceKind, EvidenceOrigin, EvidenceOriginKind, Frontmatter, SourceEvidenceRef, SourceRef, + split_frontmatter, +}; pub use decision::{DecisionFrontmatter, DecisionStatus}; pub use request::RequestFrontmatter; pub use summary::SummaryFrontmatter; diff --git a/crates/protocol/src/lib.rs b/crates/protocol/src/lib.rs index 5d7d82bf..bfc1a49a 100644 --- a/crates/protocol/src/lib.rs +++ b/crates/protocol/src/lib.rs @@ -352,6 +352,18 @@ pub struct InternalWorkerSnapshot { pub internal_workers: Vec, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] +#[serde(rename_all = "snake_case")] +pub enum ToolResultDisposition { + #[default] + Success, + Error, + Interrupted, + Cancelled, + OutcomeUnknown, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[serde(tag = "event", content = "data", rename_all = "snake_case")] @@ -501,6 +513,8 @@ pub enum Event { /// summary-only, or when the result was pruned. #[serde(default, skip_serializing_if = "Option::is_none")] output: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + disposition: Option, #[serde(default)] is_error: bool, }, @@ -923,6 +937,7 @@ pub enum WorkerStatus { Idle, Running, Paused, + Stopped, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -1838,6 +1853,7 @@ mod tests { id: "call_1".into(), summary: "Read 128 bytes".into(), output: Some("hello world".into()), + disposition: Some(ToolResultDisposition::Success), is_error: false, }; let json = serde_json::to_string(&event).unwrap(); @@ -1854,11 +1870,13 @@ mod tests { id, summary, output, + disposition, is_error, } => { assert_eq!(id, "call_1"); assert_eq!(summary, "Read 128 bytes"); assert_eq!(output.as_deref(), Some("hello world")); + assert_eq!(disposition, Some(ToolResultDisposition::Success)); assert!(!is_error); } other => panic!("expected ToolResult, got {other:?}"), @@ -1871,6 +1889,7 @@ mod tests { id: "call_2".into(), summary: "ok".into(), output: None, + disposition: Some(ToolResultDisposition::Success), is_error: false, }; let json = serde_json::to_string(&event).unwrap(); @@ -1886,6 +1905,7 @@ mod tests { id: "call_3".into(), summary: "invalid argument".into(), output: None, + disposition: Some(ToolResultDisposition::Error), is_error: true, }; let json = serde_json::to_string(&event).unwrap(); diff --git a/crates/protocol/src/typescript.rs b/crates/protocol/src/typescript.rs index 29dffce8..2ac5446c 100644 --- a/crates/protocol/src/typescript.rs +++ b/crates/protocol/src/typescript.rs @@ -8,7 +8,7 @@ use crate::{ CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot, InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot, InvokeKind, MemoryWorkerEvent, Method, Permission, RewindSummary, RewindTarget, RewindTargetId, - RunResult, ScopeRule, Segment, TurnResult, WorkerEvent, WorkerStatus, + RunResult, ScopeRule, Segment, ToolResultDisposition, TurnResult, WorkerEvent, WorkerStatus, subscription::{ EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame, SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest, @@ -45,6 +45,7 @@ pub fn generated_protocol_types() -> String { push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); diff --git a/crates/session-store/src/history.rs b/crates/session-store/src/history.rs new file mode 100644 index 00000000..3903213c --- /dev/null +++ b/crates/session-store/src/history.rs @@ -0,0 +1,180 @@ +//! Serializable history entries with restore-authoritative logical identity and origin. + +use serde::{Deserialize, Serialize}; + +use crate::{LoggedItem, SessionId}; + +/// Stable logical identity of one model-visible history entry. +/// +/// This value is generated at the trusted Worker session boundary and copied +/// unchanged across fork, rewind, compaction retention, restore, and reboot. +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct LoggedSessionHistoryEntryId(pub String); + +impl LoggedSessionHistoryEntryId { + pub fn new() -> Self { + Self(uuid::Uuid::now_v7().to_string()) + } +} + +impl Default for LoggedSessionHistoryEntryId { + fn default() -> Self { + Self::new() + } +} + +/// Bounded subject snapshot. It is evidence, not a live authorization handle. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct LoggedWorkerSubject { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub workspace_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub runtime_id: Option, + pub worker_id: String, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum LoggedSessionHistoryOrigin { + HumanInput { + account_id: String, + }, + WorkerInput { + actor: LoggedWorkerSubject, + }, + FlowInstruction { + selector: String, + definition_id: String, + definition_revision: u64, + instance_id: String, + state_id: String, + }, + BackendInstruction { + #[serde(default, skip_serializing_if = "Option::is_none")] + operation_id: Option, + }, + ModelOutput { + worker: LoggedWorkerSubject, + }, + ToolOutput { + worker: LoggedWorkerSubject, + }, + DerivedSummary, + LegacyUnknown, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct LoggedHistoryDerivation { + pub sources: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct LoggedSessionHistoryMetadata { + pub entry_id: LoggedSessionHistoryEntryId, + pub origin: LoggedSessionHistoryOrigin, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub derivation: Option, +} + +impl LoggedSessionHistoryMetadata { + pub fn legacy_unknown() -> Self { + Self { + entry_id: LoggedSessionHistoryEntryId::new(), + origin: LoggedSessionHistoryOrigin::LegacyUnknown, + derivation: None, + } + } +} + +/// Persisted item and metadata are one value so transforms cannot reorder or +/// truncate one without the other. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct LoggedHistoryEntry { + pub item: LoggedItem, + pub metadata: LoggedSessionHistoryMetadata, +} + +/// Typed system-item history record. The typed system event remains available +/// to client replay while its model-visible projection carries the same stable +/// metadata used by live history. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct LoggedSystemHistoryEntry { + pub item: crate::SystemItem, + pub metadata: LoggedSessionHistoryMetadata, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::LoggedRole; + use agen::llm_client::RequestConfig; + + #[test] + fn logged_history_entry_round_trip_preserves_id_origin_and_derivation() { + let source_id = LoggedSessionHistoryEntryId::new(); + let entry = LoggedHistoryEntry { + item: LoggedItem::Message { + role: LoggedRole::User, + content: vec![crate::LoggedContentPart::Text { + text: "preference".into(), + }], + }, + metadata: LoggedSessionHistoryMetadata { + entry_id: LoggedSessionHistoryEntryId::new(), + origin: LoggedSessionHistoryOrigin::HumanInput { + account_id: "account-1".into(), + }, + derivation: Some(LoggedHistoryDerivation { + sources: vec![source_id.clone()], + }), + }, + }; + let encoded = serde_json::to_vec(&entry).unwrap(); + let decoded: LoggedHistoryEntry = serde_json::from_slice(&encoded).unwrap(); + assert_eq!(decoded, entry); + assert_eq!( + decoded.metadata.derivation.unwrap().sources, + vec![source_id] + ); + } + + #[test] + fn annotated_segment_start_is_restore_visible_without_projecting_metadata() { + let session_id = uuid::Uuid::now_v7(); + let history_entry = legacy_logged_history(LoggedItem::Message { + role: LoggedRole::Assistant, + content: vec![crate::LoggedContentPart::Text { + text: "answer".into(), + }], + }); + let state = crate::collect_state(&[crate::LogEntry::AnnotatedSegmentStart { + ts: 1, + session_id, + system_prompt: None, + config: RequestConfig::default(), + history: vec![history_entry], + forked_from: None, + compacted_from: None, + }]); + assert_eq!(state.history[0].as_text(), Some("answer")); + } +} + +/// Legacy Session Logs did not persist annotations. Decode helpers explicitly +/// create `LegacyUnknown`; they never infer Human/System authority from role or +/// plaintext. +pub fn legacy_logged_history(item: LoggedItem) -> LoggedHistoryEntry { + LoggedHistoryEntry { + item, + metadata: LoggedSessionHistoryMetadata::legacy_unknown(), + } +} + +pub fn legacy_segment_history( + session_id: SessionId, + items: impl IntoIterator, +) -> Vec { + let _ = session_id; + items.into_iter().map(legacy_logged_history).collect() +} diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index 28f5c708..9fe1fad7 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -32,6 +32,7 @@ pub mod event_trace; pub mod fs_store; +pub mod history; pub mod logged_item; pub mod segment; pub mod segment_log; @@ -44,6 +45,11 @@ pub use agen::UsageRecord; pub use agen::llm_client::types::{ContentPart, Item, Role}; pub use event_trace::{TraceEntry, TracePayload}; pub use fs_store::FsStore; +pub use history::{ + LoggedHistoryDerivation, LoggedHistoryEntry, LoggedSessionHistoryEntryId, + LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, LoggedSystemHistoryEntry, + LoggedWorkerSubject, legacy_logged_history, legacy_segment_history, +}; pub use logged_item::{LoggedContentPart, LoggedItem, LoggedRole, from_logged, to_logged}; pub use segment::{ SegmentStartState, append_entry, append_system_item, classify_history_item, diff --git a/crates/session-store/src/logged_item.rs b/crates/session-store/src/logged_item.rs index f2d499e6..f6547265 100644 --- a/crates/session-store/src/logged_item.rs +++ b/crates/session-store/src/logged_item.rs @@ -14,7 +14,7 @@ use agen::{ llm_client::types::{ContentPart, Item, Role}, - tool::{Attachment, ImageAttachment}, + tool::{Attachment, ImageAttachment, ToolResultDisposition}, }; use base64::{Engine as _, engine::general_purpose::STANDARD}; use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _}; @@ -61,6 +61,8 @@ pub enum LoggedItem { content: Option, #[serde(default, skip_serializing_if = "Vec::is_empty")] attachments: Vec, + #[serde(default, skip_serializing_if = "ToolResultDisposition::is_success")] + disposition: ToolResultDisposition, #[serde(default, skip_serializing_if = "is_false")] is_error: bool, }, @@ -128,6 +130,7 @@ impl From<&Item> for LoggedItem { summary, content, attachments, + disposition, is_error, .. } => Self::ToolResult { @@ -135,6 +138,7 @@ impl From<&Item> for LoggedItem { summary: summary.clone(), content: content.clone(), attachments: attachments.iter().map(LoggedAttachment::from).collect(), + disposition: *disposition, is_error: *is_error, }, Item::Reasoning { @@ -184,15 +188,24 @@ impl From for Item { summary, content, attachments, + disposition, is_error, - } => Item::ToolResult { - id: None, - call_id, - summary, - content, - is_error, - attachments: attachments.into_iter().map(Attachment::from).collect(), - }, + } => { + let disposition = if is_error && disposition.is_success() { + ToolResultDisposition::Error + } else { + disposition + }; + Item::ToolResult { + id: None, + call_id, + summary, + content, + disposition, + is_error, + attachments: attachments.into_iter().map(Attachment::from).collect(), + } + } LoggedItem::Reasoning { text, summary, @@ -430,6 +443,42 @@ mod tests { } } + #[test] + fn outcome_unknown_tool_result_round_trips_as_terminal() { + let original = Item::tool_result_item_with_disposition_and_attachments( + "call_unknown", + "outcome unknown", + Some("bounded progress".to_string()), + ToolResultDisposition::OutcomeUnknown, + Vec::new(), + ); + let logged: LoggedItem = (&original).into(); + let json = serde_json::to_string(&logged).unwrap(); + assert!(json.contains(r#""disposition":"outcome_unknown""#)); + match Item::from(serde_json::from_str::(&json).unwrap()) { + Item::ToolResult { + disposition, + is_error, + .. + } => { + assert_eq!(disposition, ToolResultDisposition::OutcomeUnknown); + assert!(is_error); + } + other => panic!("unexpected variant: {other:?}"), + } + } + + #[test] + fn legacy_error_tool_result_infers_error_disposition() { + let legacy = r#"{"kind":"tool_result","call_id":"call_old","summary":"failed","content":null,"is_error":true}"#; + match Item::from(serde_json::from_str::(legacy).unwrap()) { + Item::ToolResult { disposition, .. } => { + assert_eq!(disposition, ToolResultDisposition::Error) + } + other => panic!("unexpected variant: {other:?}"), + } + } + #[test] fn tool_result_persistence_round_trips_binary_attachments() { let original = Item::tool_result_item_with_attachments( diff --git a/crates/session-store/src/segment_log.rs b/crates/session-store/src/segment_log.rs index 5bf47454..b03163d2 100644 --- a/crates/session-store/src/segment_log.rs +++ b/crates/session-store/src/segment_log.rs @@ -14,6 +14,7 @@ use agen::{EngineResult, UsageRecord}; use protocol::{InvokeKind, Segment}; use serde::{Deserialize, Serialize}; +use crate::history::{LoggedHistoryEntry, LoggedSystemHistoryEntry}; use crate::logged_item::LoggedItem; use crate::system_item::SystemItem; @@ -70,6 +71,20 @@ pub enum LogEntry { compacted_from: Option, }, + /// Schema-v2 segment seed. Retained entries keep their stable logical + /// identity and origin across fork/compaction/restore. + AnnotatedSegmentStart { + ts: u64, + session_id: crate::SessionId, + system_prompt: Option, + config: RequestConfig, + history: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + forked_from: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + compacted_from: Option, + }, + /// IDLE → active marker. Records the start of a new self-driving /// cycle (Invoke range). The range extends implicitly until the /// next `Invoke` entry; this entry carries the trigger only — the @@ -105,14 +120,37 @@ pub enum LogEntry { extensions: Vec, }, + /// Schema-v2 user submission with its exact model-visible entries. Typed + /// Flow instructions and caller-attributed input remain separate entries. + AnnotatedUserInput { + ts: u64, + segments: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + extensions: Vec, + history: Vec, + }, + + /// Schema-v2 model output and metadata committed as one journal record. + AnnotatedAssistantItem { ts: u64, entry: LoggedHistoryEntry }, + /// One assistant-side item appended to history — assistant message, /// reasoning, or tool call. Singular: one entry per history item so /// the wire-side `Event::*` lane and on-disk LogEntry stay 1:1. AssistantItem { ts: u64, item: LoggedItem }, + /// Schema-v2 tool output and metadata committed as one journal record. + AnnotatedToolResult { ts: u64, entry: LoggedHistoryEntry }, + /// One tool-execution result appended to history. ToolResult { ts: u64, item: LoggedItem }, + /// Schema-v2 typed system event and model-visible metadata committed + /// together. + AnnotatedSystemItem { + ts: u64, + entry: LoggedSystemHistoryEntry, + }, + /// One typed agent-injected system item: notification, child-Worker /// lifecycle event, `@` / `/` resolution payload. Each /// `SystemItem` carries kind metadata that the LLM @@ -278,6 +316,22 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState { state.config = config.clone(); state.history = history.iter().cloned().map(Item::from).collect(); } + LogEntry::AnnotatedSegmentStart { + session_id, + system_prompt, + config, + history, + .. + } => { + state.session_id = Some(*session_id); + state.system_prompt = system_prompt.clone(); + state.config = config.clone(); + state.history = history + .iter() + .cloned() + .map(|entry| Item::from(entry.item)) + .collect(); + } LogEntry::Invoke { .. } => { // A terminal run record below clears or refines this. If the // log ends first, restore must treat the turn as interrupted. @@ -298,6 +352,29 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState { .map(|extension| (extension.domain.clone(), extension.payload.clone())), ); } + LogEntry::AnnotatedUserInput { + segments, + extensions, + history, + .. + } => { + state + .history + .extend(history.iter().cloned().map(|entry| Item::from(entry.item))); + state.user_segments.push(segments.clone()); + state.extensions.extend( + extensions + .iter() + .map(|extension| (extension.domain.clone(), extension.payload.clone())), + ); + } + LogEntry::AnnotatedAssistantItem { entry, .. } + | LogEntry::AnnotatedToolResult { entry, .. } => { + state.history.push(Item::from(entry.item.clone())); + } + LogEntry::AnnotatedSystemItem { entry, .. } => { + state.history.push(entry.item.to_history_item()); + } LogEntry::AssistantItem { item, .. } => { state.history.push(Item::from(item.clone())); } diff --git a/crates/session-store/src/worker_session_store.rs b/crates/session-store/src/worker_session_store.rs index 902b5dd8..5dd063d8 100644 --- a/crates/session-store/src/worker_session_store.rs +++ b/crates/session-store/src/worker_session_store.rs @@ -20,7 +20,8 @@ use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; use std::time::SystemTime; -const SESSION_SCHEMA_VERSION: u32 = 1; +const SESSION_SCHEMA_VERSION: u32 = 2; +const LEGACY_SESSION_SCHEMA_VERSION: u32 = 1; const SESSION_FILE: &str = "session.json"; const SEGMENTS_DIR: &str = "segments"; @@ -44,15 +45,22 @@ impl WorkerSessionStore { fs::create_dir_all(root.join(SEGMENTS_DIR))?; let session_id = match fs::read(root.join(SESSION_FILE)) { Ok(bytes) => { - let manifest: SessionManifest = serde_json::from_slice(&bytes)?; - if manifest.schema_version != SESSION_SCHEMA_VERSION { - return Err(StoreError::Corrupt { - line: 0, - message: format!( - "unsupported Worker Session schema version {}, expected {}", - manifest.schema_version, SESSION_SCHEMA_VERSION - ), - }); + let mut manifest: SessionManifest = serde_json::from_slice(&bytes)?; + match manifest.schema_version { + SESSION_SCHEMA_VERSION => {} + LEGACY_SESSION_SCHEMA_VERSION => { + validate_legacy_segment_logs(&root)?; + manifest.schema_version = SESSION_SCHEMA_VERSION; + atomic_write_json(&root.join(SESSION_FILE), &manifest)?; + } + version => { + return Err(StoreError::Corrupt { + line: 0, + message: format!( + "unsupported Worker Session schema version {version}, expected {SESSION_SCHEMA_VERSION}" + ), + }); + } } Some(manifest.session_id) } @@ -278,6 +286,37 @@ impl Store for WorkerSessionStore { } } +fn validate_legacy_segment_logs(root: &Path) -> Result<(), StoreError> { + let segments = root.join(SEGMENTS_DIR); + if !segments.exists() { + return Ok(()); + } + for entry in fs::read_dir(&segments)? { + let entry = entry?; + let path = entry.path(); + let Some(name) = path.file_name().and_then(|name| name.to_str()) else { + continue; + }; + if !name.ends_with(".jsonl") || name.ends_with(".trace.jsonl") { + continue; + } + let contents = fs::read_to_string(&path)?; + for (line_index, line) in contents.lines().enumerate() { + if line.trim().is_empty() { + continue; + } + serde_json::from_str::(line).map_err(|error| StoreError::Corrupt { + line: line_index + 1, + message: format!( + "cannot migrate legacy Worker Session log {}: {error}", + path.display() + ), + })?; + } + } + Ok(()) +} + fn atomic_write_json(path: &Path, value: &T) -> Result<(), StoreError> { let mut bytes = serde_json::to_vec_pretty(value)?; bytes.push(b'\n'); @@ -405,6 +444,54 @@ mod tests { assert_eq!(store.list_sessions().unwrap(), vec![session_id]); } + #[test] + fn schema_v1_logs_are_validated_and_promoted_to_v2() { + let root = tempfile::tempdir().unwrap(); + let session_id = new_session_id(); + let segment_id = new_segment_id(); + WorkerSessionStore::new(root.path()) + .unwrap() + .create_segment(session_id, segment_id, &[]) + .unwrap(); + let manifest_path = root.path().join(SESSION_FILE); + let mut manifest: SessionManifest = + serde_json::from_slice(&fs::read(&manifest_path).unwrap()).unwrap(); + manifest.schema_version = LEGACY_SESSION_SCHEMA_VERSION; + atomic_write_json(&manifest_path, &manifest).unwrap(); + + let reopened = WorkerSessionStore::new(root.path()).unwrap(); + assert_eq!(reopened.session_id().unwrap(), Some(session_id)); + let migrated: SessionManifest = + serde_json::from_slice(&fs::read(&manifest_path).unwrap()).unwrap(); + assert_eq!(migrated.schema_version, SESSION_SCHEMA_VERSION); + } + + #[test] + fn schema_v1_migration_rejects_corrupt_log_before_manifest_update() { + let root = tempfile::tempdir().unwrap(); + let session_id = new_session_id(); + let manifest = SessionManifest { + schema_version: LEGACY_SESSION_SCHEMA_VERSION, + session_id, + }; + atomic_write_json(&root.path().join(SESSION_FILE), &manifest).unwrap(); + fs::create_dir_all(root.path().join(SEGMENTS_DIR)).unwrap(); + fs::write( + root.path().join(SEGMENTS_DIR).join("broken.jsonl"), + "{not-json}\n", + ) + .unwrap(); + + let error = match WorkerSessionStore::new(root.path()) { + Ok(_) => panic!("corrupt legacy Session log must reject migration"), + Err(error) => error, + }; + assert!(matches!(error, StoreError::Corrupt { .. })); + let persisted: SessionManifest = + serde_json::from_slice(&fs::read(root.path().join(SESSION_FILE)).unwrap()).unwrap(); + assert_eq!(persisted.schema_version, LEGACY_SESSION_SCHEMA_VERSION); + } + #[test] fn reopen_preserves_session_and_segment_ids() { let root = tempfile::tempdir().unwrap(); diff --git a/crates/session-store/tests/session_test.rs b/crates/session-store/tests/session_test.rs index dda7205d..d8e27fe0 100644 --- a/crates/session-store/tests/session_test.rs +++ b/crates/session-store/tests/session_test.rs @@ -1,12 +1,13 @@ mod common; +use std::ops::{Deref, DerefMut}; use std::sync::Arc; -use agen::Engine; use agen::interceptor::{Interceptor, TurnEndAction}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::types::{Item, RequestConfig}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; +use agen::{Engine, History}; use async_trait::async_trait; use common::MockLlmClient; use session_store::{FsStore, LogEntry, SegmentStartState, Store, collect_state}; @@ -94,15 +95,47 @@ fn make_store() -> (tempfile::TempDir, FsStore) { (dir, store) } +struct TestWorker { + engine: Engine, + history: History, +} + +impl TestWorker { + fn new(engine: Engine) -> Self { + Self { + engine, + history: History::new(), + } + } + + fn history(&self) -> Vec { + self.history.items_cloned() + } +} + +impl Deref for TestWorker { + type Target = Engine; + + fn deref(&self) -> &Self::Target { + &self.engine + } +} + +impl DerefMut for TestWorker { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.engine + } +} + /// Run a worker turn and persist via session-store functions. /// Takes ownership of the worker (needed for lock/unlock) and returns it. async fn run_and_persist( - worker: Engine, + mut worker: TestWorker, store: &FsStore, session_id: session_store::SessionId, segment_id: session_store::SegmentId, input: &str, -) -> (Engine, agen::EngineResult) { +) -> (TestWorker, agen::EngineRunExit) { // Mirror Worker's run-entry contract: log the user input as segments // before the worker pushes its flattened user_message; save_delta // skips the resulting user_message item to avoid double-write. @@ -114,42 +147,61 @@ async fn run_and_persist( ) .unwrap(); - let history_before = worker.history().len(); + let history_before = worker.history.len(); - let mut locked = worker.lock(); - let result = locked.run(input).await; - let worker = locked.unlock(); + let mut locked = worker.engine.lock(&worker.history); + let result = locked.run(&mut worker.history, input).await; + worker.engine = locked.unlock(); - let new_items = &worker.history()[history_before..]; + let projected = worker.history(); + let new_items = &projected[history_before..]; session_store::save_delta(store, session_id, segment_id, new_items).unwrap(); session_store::save_turn_end(store, session_id, segment_id, worker.turn_count()).unwrap(); match &result { - Ok(r) => { + agen::EngineRunExit::Finished + | agen::EngineRunExit::Paused + | agen::EngineRunExit::Yielded => { + let (legacy_result, interrupted) = match &result { + agen::EngineRunExit::Finished => (agen::EngineResult::Finished, false), + agen::EngineRunExit::Paused => (agen::EngineResult::Paused, true), + agen::EngineRunExit::Yielded => (agen::EngineResult::Yielded, true), + agen::EngineRunExit::Interrupted(_) => unreachable!(), + }; session_store::save_run_completed( store, session_id, segment_id, - r.clone(), - worker.last_run_interrupted(), + legacy_result, + interrupted, worker.active_run_turn_count(), ) .unwrap(); } - Err(e) => { + agen::EngineRunExit::Interrupted(agen::StopReason::LimitReached) => { + session_store::save_run_completed( + store, + session_id, + segment_id, + agen::EngineResult::LimitReached, + false, + worker.active_run_turn_count(), + ) + .unwrap(); + } + agen::EngineRunExit::Interrupted(reason) => { session_store::save_run_errored( store, session_id, segment_id, - e.to_string(), - worker.last_run_interrupted(), + format!("{reason:?}"), + true, ) .unwrap(); } } - let r = result.unwrap(); - (worker, r) + (worker, result) } // ============================================================================= @@ -160,14 +212,14 @@ async fn run_and_persist( async fn session_run_logs_entries() { let (_dir, store) = make_store(); let client = MockLlmClient::new(simple_text_events()); - let worker = Engine::new(client); + let worker = TestWorker::new(Engine::new(client)); let (sid, segid) = session_store::create_segment( &store, SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: worker.history(), + history: &worker.history(), }, ) .unwrap(); @@ -204,7 +256,7 @@ async fn session_run_logs_entries() { async fn session_restore_round_trip() { let (_dir, store) = make_store(); let client = MockLlmClient::new(simple_text_events()); - let mut worker = Engine::new(client); + let mut worker = TestWorker::new(Engine::new(client)); worker.set_system_prompt("You are helpful."); let (sid, segid) = session_store::create_segment( @@ -212,7 +264,7 @@ async fn session_restore_round_trip() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: worker.history(), + history: &worker.history(), }, ) .unwrap(); @@ -243,7 +295,7 @@ async fn session_restore_round_trip() { async fn session_run_with_tool_call() { let (_dir, store) = make_store(); let client = MockLlmClient::with_responses(tool_call_events()); - let mut worker = Engine::new(client); + let mut worker = TestWorker::new(Engine::new(client)); worker.register_tool(weather_tool_definition()); let (sid, segid) = session_store::create_segment( @@ -251,7 +303,7 @@ async fn session_run_with_tool_call() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: worker.history(), + history: &worker.history(), }, ) .unwrap(); @@ -277,7 +329,7 @@ async fn session_resume_after_pause() { // First run: tool call with pause policy → Paused let client = MockLlmClient::with_responses(tool_call_events()); - let mut worker = Engine::new(client); + let mut worker = TestWorker::new(Engine::new(client)); worker.register_tool(weather_tool_definition()); worker.set_interceptor(PausePolicy); @@ -286,13 +338,13 @@ async fn session_resume_after_pause() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: worker.history(), + history: &worker.history(), }, ) .unwrap(); let (_worker, result) = run_and_persist(worker, &store, sid, segid, "Weather?").await; - assert!(matches!(result, agen::EngineResult::Paused)); + assert!(matches!(result, agen::EngineRunExit::Paused)); // Check RunCompleted is Paused let entries = store.read_all(sid, segid).unwrap(); @@ -317,7 +369,7 @@ async fn session_resume_after_pause() { async fn session_fork_creates_new_session() { let (_dir, store) = make_store(); let client = MockLlmClient::new(simple_text_events()); - let mut worker = Engine::new(client); + let mut worker = TestWorker::new(Engine::new(client)); worker.set_system_prompt("System prompt"); let (sid, segid) = session_store::create_segment( @@ -325,7 +377,7 @@ async fn session_fork_creates_new_session() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: worker.history(), + history: &worker.history(), }, ) .unwrap(); @@ -338,7 +390,7 @@ async fn session_fork_creates_new_session() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: worker.history(), + history: &worker.history(), }, ) .unwrap(); @@ -359,14 +411,14 @@ async fn session_fork_creates_new_session() { async fn session_fork_at_truncates_within_session() { let (_dir, store) = make_store(); let client = MockLlmClient::new(simple_text_events()); - let worker = Engine::new(client); + let worker = TestWorker::new(Engine::new(client)); let (sid, segid) = session_store::create_segment( &store, SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: worker.history(), + history: &worker.history(), }, ) .unwrap(); @@ -404,14 +456,14 @@ async fn session_fork_at_truncates_within_session() { async fn session_config_changed_logged() { let (_dir, store) = make_store(); let client = MockLlmClient::new(vec![]); - let mut worker = Engine::new(client); + let mut worker = TestWorker::new(Engine::new(client)); let (sid, segid) = session_store::create_segment( &store, SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: worker.history(), + history: &worker.history(), }, ) .unwrap(); @@ -437,14 +489,14 @@ async fn session_auto_forks_on_conflict() { // Create a segment let client_a = MockLlmClient::new(simple_text_events()); - let worker_a = Engine::new(client_a); + let worker_a = TestWorker::new(Engine::new(client_a)); let (sid, original_segid) = session_store::create_segment( &store, SegmentStartState { system_prompt: worker_a.get_system_prompt(), config: worker_a.request_config(), - history: worker_a.history(), + history: &worker_a.history(), }, ) .unwrap(); @@ -470,7 +522,7 @@ async fn session_auto_forks_on_conflict() { SegmentStartState { system_prompt: worker_a.get_system_prompt(), config: worker_a.request_config(), - history: worker_a.history(), + history: &worker_a.history(), }, ) .unwrap(); @@ -522,14 +574,14 @@ async fn session_auto_forks_on_conflict() { async fn nested_past_fork_leaves_ancestors_immutable() { let (_dir, store) = make_store(); let client = MockLlmClient::new(simple_text_events()); - let worker = Engine::new(client); + let worker = TestWorker::new(Engine::new(client)); let (sid, root_segid) = session_store::create_segment( &store, SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: worker.history(), + history: &worker.history(), }, ) .unwrap(); diff --git a/crates/ticket/src/tool.rs b/crates/ticket/src/tool.rs index f6fb2842..a38615b2 100644 --- a/crates/ticket/src/tool.rs +++ b/crates/ticket/src/tool.rs @@ -410,7 +410,7 @@ struct TicketCreateParams { #[derive(Debug, Deserialize, schemars::JsonSchema)] struct TicketEditItemParams { - /// Ticket id. + /// Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. ticket: String, /// Optional replacement title. #[serde(default)] @@ -539,7 +539,7 @@ impl QueryTicketParams { #[derive(Debug, Deserialize, schemars::JsonSchema)] struct ShowTicketParams { - /// Ticket id. Exactly one of `id` or `query` must be provided. + /// Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. Exactly one of `id` or `query` must be provided. #[serde(default)] id: Option, /// Exact ticket id query. Exactly one of `id` or `query` must be provided. @@ -558,7 +558,7 @@ struct ShowTicketParams { #[derive(Debug, Deserialize, schemars::JsonSchema)] struct TicketThreadEventParams { - /// Ticket id. + /// Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. ticket: String, /// Markdown event body. body: String, @@ -566,7 +566,7 @@ struct TicketThreadEventParams { #[derive(Debug, Deserialize, schemars::JsonSchema)] struct TicketMarkReadyParams { - /// Ticket id. + /// Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. ticket: String, /// Optional reason attached to the state_changed event. #[serde(default)] @@ -575,7 +575,7 @@ struct TicketMarkReadyParams { #[derive(Debug, Deserialize, schemars::JsonSchema)] struct TicketIntakeReadyParams { - /// Ticket id. + /// Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. ticket: String, /// Concise bounded intake summary appended before the ready transition. intake_summary: String, @@ -586,13 +586,13 @@ struct TicketIntakeReadyParams { #[derive(Debug, Deserialize, schemars::JsonSchema)] struct TicketQueueParams { - /// Ticket id. + /// Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. ticket: String, } #[derive(Debug, Deserialize, schemars::JsonSchema)] struct TicketWorkflowStateParams { - /// Ticket id. + /// Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. ticket: String, /// Expected current state. The backend rejects stale transitions. from: TicketWorkflowStateParam, @@ -606,7 +606,7 @@ struct TicketWorkflowStateParams { #[derive(Debug, Deserialize, schemars::JsonSchema)] struct TicketCloseParams { - /// Ticket id. + /// Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. ticket: String, /// Markdown resolution written to resolution.md and thread.md. resolution: String, @@ -614,7 +614,7 @@ struct TicketCloseParams { #[derive(Debug, Deserialize, schemars::JsonSchema)] struct TicketDependencyCheckParams { - /// Ticket id. + /// Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. ticket: String, } @@ -646,7 +646,7 @@ struct TicketRelationRecordParams { ticket: String, /// Forward relation kind: depends_on, blocks, related, supersedes, or duplicate_of. kind: TicketRelationKindParam, - /// Target canonical Ticket id. Title/slug words are not accepted as relation authority. + /// Target Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. target: String, /// Optional bounded rationale/note. #[serde(default)] @@ -659,7 +659,7 @@ struct TicketRelationRemoveParams { ticket: String, /// Forward relation kind to remove. kind: TicketRelationKindParam, - /// Target canonical Ticket id. + /// Target Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. target: String, } @@ -1223,10 +1223,17 @@ impl Tool for TicketQueueTool { ) -> Result { let params: TicketQueueParams = parse_input("TicketQueue", input_json)?; let queued_by = default_author(); - let outcome = self + let mut outcome = self .backend .queue_ready(TicketIdOrSlug::Query(params.ticket.clone()), &queued_by) .map_err(|error| backend_error("TicketQueue", error))?; + outcome.requested_ticket = + model_ticket_reference(&self.backend, &outcome.requested_ticket, "TicketQueue")?; + outcome.queued_tickets = outcome + .queued_tickets + .into_iter() + .map(|ticket| model_ticket_reference(&self.backend, &ticket, "TicketQueue")) + .collect::, _>>()?; Ok(json_output( format!( "Queued {} ticket(s) for Orchestrator", @@ -1264,15 +1271,17 @@ impl Tool for TicketWorkflowStateTool { self.backend .set_workflow_state(TicketIdOrSlug::Query(params.ticket.clone()), change) .map_err(|error| backend_error("TicketWorkflowState", error))?; + let ticket_ref = + model_ticket_reference(&self.backend, ¶ms.ticket, "TicketWorkflowState")?; Ok(json_output( format!( "Transitioned ticket {} state {} -> {}", - params.ticket, + ticket_ref, from.as_str(), to.as_str() ), json!({ - "ticket": params.ticket, + "ticket": ticket_ref, "from": from.as_str(), "to": to.as_str(), "state": to.as_str(), @@ -1296,9 +1305,10 @@ impl Tool for TicketCloseTool { MarkdownText::new(params.resolution), ) .map_err(|error| backend_error("TicketClose", error))?; + let ticket_ref = model_ticket_reference(&self.backend, ¶ms.ticket, "TicketClose")?; Ok(json_output( - format!("Closed ticket {}", params.ticket), - json!({ "ticket": params.ticket, "state": "closed", "ok": true }), + format!("Closed ticket {ticket_ref}"), + json!({ "ticket": ticket_ref, "state": "closed", "ok": true }), )) } } @@ -1525,6 +1535,29 @@ impl Tool for TicketDependencyCheckTool { } } +fn model_ticket_reference( + backend: &TicketToolBackend, + reference: &str, + tool_name: &str, +) -> Result { + let ticket = backend + .show(TicketIdOrSlug::Id(reference.to_string())) + .map_err(|error| backend_error(tool_name, error))?; + match ticket.meta.resource_key { + Some(resource_key) if is_canonical_ticket_resource_key(&resource_key) => Ok(resource_key), + Some(_) => Err(ToolError::ExecutionFailed(format!( + "{tool_name} failed: required Ticket human key is unavailable" + ))), + None => Ok(ticket.meta.id), + } +} + +fn is_canonical_ticket_resource_key(resource_key: &str) -> bool { + resource_key.strip_prefix("T-").is_some_and(|sequence| { + !sequence.is_empty() && sequence.bytes().all(|byte| byte.is_ascii_digit()) + }) +} + fn parse_input Deserialize<'de>>(tool: &str, input_json: &str) -> Result { serde_json::from_str(input_json) .map_err(|error| ToolError::InvalidArgument(format!("invalid {tool} input: {error}"))) @@ -1922,6 +1955,12 @@ mod tests { .with_target_authority(Arc::new(TestTargetAuthority)) } + fn sqlite_backend(temp: &TempDir) -> crate::SqliteTicketBackend { + crate::SqliteTicketBackend::open(temp.path().join("tickets.db"), "workspace") + .unwrap() + .with_target_authority(Arc::new(TestTargetAuthority)) + } + fn tool(definition: ToolDefinition) -> Arc { let (_, tool) = definition(); tool @@ -2549,6 +2588,101 @@ mod tests { ); } + #[tokio::test] + async fn queue_workflow_and_close_project_internal_inputs_to_ticket_keys() { + let temp = TempDir::new().unwrap(); + let inner = sqlite_backend(&temp); + let mut dependency_input = NewTicket::new("Dependency"); + dependency_input.repository_id = Some("main".to_string()); + let dependency = inner.create(dependency_input).unwrap(); + let mut target_input = NewTicket::new("Target"); + target_input.repository_id = Some("main".to_string()); + let target = inner.create(target_input).unwrap(); + inner + .add_ticket_relation( + TicketIdOrSlug::Id(target.id.clone()), + NewTicketRelation { + kind: TicketRelationKind::DependsOn, + target: dependency.id.clone(), + note: None, + author: None, + }, + ) + .unwrap(); + for id in [&dependency.id, &target.id] { + inner + .mark_ready( + TicketIdOrSlug::Id(id.clone()), + TicketMarkReady { + operation_key: format!("ready-{id}"), + reason: None, + author: None, + intake_summary: None, + }, + ) + .unwrap(); + } + let target_key = target.resource_key.clone().unwrap(); + let dependency_key = dependency.resource_key.clone().unwrap(); + let backend = inner; + let queue = tool_by_name(TicketToolBackend::new(backend.clone()), "TicketQueue"); + let workflow = tool_by_name( + TicketToolBackend::new(backend.clone()), + "TicketWorkflowState", + ); + let close = tool_by_name(TicketToolBackend::new(backend), "TicketClose"); + + let queued = queue + .execute( + &json!({"ticket": target.id.clone()}).to_string(), + Default::default(), + ) + .await + .unwrap(); + assert!(queued.summary.contains("2 ticket(s)")); + let queued_content = queued.content.unwrap(); + assert!(queued_content.contains(&target_key)); + assert!(queued_content.contains(&dependency_key)); + assert!(!queued_content.contains(&target.id)); + assert!(!queued_content.contains(&dependency.id)); + + for (from, to) in [("queued", "inprogress"), ("inprogress", "done")] { + let transitioned = workflow + .execute( + &json!({ + "ticket": target.id.clone(), + "from": from, + "to": to, + "reason": "test_transition", + "body": "transitioned", + "author": "tester" + }) + .to_string(), + Default::default(), + ) + .await + .unwrap(); + assert!(transitioned.summary.contains(&target_key)); + assert!(!transitioned.summary.contains(&target.id)); + let content = transitioned.content.unwrap(); + assert!(content.contains(&target_key)); + assert!(!content.contains(&target.id)); + } + + let closed = close + .execute( + &json!({"ticket": target.id.clone(), "resolution": "Done"}).to_string(), + Default::default(), + ) + .await + .unwrap(); + assert!(closed.summary.contains(&target_key)); + assert!(!closed.summary.contains(&target.id)); + let content = closed.content.unwrap(); + assert!(content.contains(&target_key)); + assert!(!content.contains(&target.id)); + } + #[tokio::test] async fn ticket_workflow_tools_mark_ready_and_transition_state() { let temp = TempDir::new().unwrap(); diff --git a/crates/tools/src/bash.rs b/crates/tools/src/bash.rs index 8224472b..bc573849 100644 --- a/crates/tools/src/bash.rs +++ b/crates/tools/src/bash.rs @@ -1,5 +1,6 @@ +use std::collections::{HashMap, HashSet}; use std::path::PathBuf; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use async_trait::async_trait; @@ -20,21 +21,65 @@ struct BashParams { pub(crate) struct BashTool { session: WorkdirSessionHandle, + state: Arc>, +} + +#[derive(Clone)] +struct ActiveCommand { + call_id: String, + execution_nonce: u64, + handle: CommandHandle, +} + +#[derive(Default)] +struct BashExecutionState { + active: HashMap, + cancellation_requested: HashSet, + legacy_cancellation_requested: HashSet, + next_execution_nonce: u64, } struct CommandGuard { session: WorkdirSessionHandle, + state: Arc>, + execution_id: String, + execution_nonce: u64, handle: Option, } impl Drop for CommandGuard { fn drop(&mut self) { - if let Some(handle) = self.handle.take() { - let workdir = self.session.clone(); - tokio::spawn(async move { - let _ = workdir.cancel_command(handle).await; - }); - } + let Some(handle) = self.handle.take() else { + return; + }; + let workdir = self.session.clone(); + let state = Arc::clone(&self.state); + let execution_id = self.execution_id.clone(); + let execution_nonce = self.execution_nonce; + // A dropped provider future is not terminal confirmation. Keep the live + // execution registered until cleanup has both requested cancellation and + // observed terminal command output, so cancellation/session teardown + // cannot race with an apparently empty registry. + tokio::spawn(async move { + let _ = workdir.cancel_command(handle.clone()).await; + let _ = workdir + .command_output(CommandOutputRequest { + handle, + cursor: 0, + limit: INLINE_BYTE_BUDGET, + wait: true, + }) + .await; + let mut state = state.lock().unwrap(); + if state + .active + .get(&execution_id) + .is_some_and(|active| active.execution_nonce == execution_nonce) + { + state.active.remove(&execution_id); + state.cancellation_requested.remove(&execution_id); + } + }); } } @@ -52,20 +97,50 @@ impl Tool for BashTool { .unwrap_or(DEFAULT_TIMEOUT_SECS) .clamp(1, MAX_TIMEOUT_SECS); let cmd_summary = truncate_for_summary(¶ms.command); + let execution_id = ctx.execution_id(); + let call_id = ctx.call_id; + let execution_nonce = { + let mut state = self.state.lock().unwrap(); + state.next_execution_nonce = state.next_execution_nonce.wrapping_add(1); + state.next_execution_nonce + }; + let mut guard = CommandGuard { + session: self.session.clone(), + state: self.state.clone(), + execution_id: execution_id.clone(), + execution_nonce, + handle: None, + }; let handle = self .session .start_command(CommandRequest { command: params.command, timeout_secs, output_limit: INLINE_BYTE_BUDGET, - tool_call_id: Some(ctx.call_id), + tool_call_id: Some(call_id.clone()), }) .await .map_err(crate::ToolsError::from)?; - let mut guard = CommandGuard { - session: self.session.clone(), - handle: Some(handle.clone()), + let cancel_after_start = { + let mut state = self.state.lock().unwrap(); + state.active.insert( + execution_id.clone(), + ActiveCommand { + call_id: call_id.clone(), + execution_nonce, + handle: handle.clone(), + }, + ); + state.cancellation_requested.contains(&execution_id) + || state.legacy_cancellation_requested.contains(&call_id) }; + guard.handle = Some(handle.clone()); + if cancel_after_start { + self.session + .cancel_command(handle.clone()) + .await + .map_err(crate::ToolsError::from)?; + } let output = self .session .command_output(CommandOutputRequest { @@ -76,9 +151,27 @@ impl Tool for BashTool { }) .await .map_err(crate::ToolsError::from)?; + let cancellation_requested = { + let mut state = self.state.lock().unwrap(); + let owns_registration = state + .active + .get(&execution_id) + .is_some_and(|active| active.execution_nonce == execution_nonce); + let exact = if owns_registration { + state.active.remove(&execution_id); + state.cancellation_requested.remove(&execution_id) + } else { + false + }; + let legacy = state.legacy_cancellation_requested.remove(&call_id); + exact || legacy + }; guard.handle = None; - let summary = if output.timed_out { + let timed_out = output.timed_out; + let summary = if cancellation_requested { + format!("$ {cmd_summary} (cancelled)") + } else if output.timed_out { format!("$ {cmd_summary} (timed out after {timeout_secs}s)") } else { match output.exit_code { @@ -97,11 +190,62 @@ impl Tool for BashTool { } else { Some(output.content) }; - Ok(ToolOutput { + let output = ToolOutput { summary, content, attachments: Vec::new(), - }) + }; + if cancellation_requested { + Err(ToolError::Cancelled(output)) + } else if timed_out { + Err(ToolError::Interrupted(output)) + } else { + Ok(output) + } + } + + async fn cancel(&self, call_id: &str) -> Result<(), ToolError> { + let handles = { + let mut state = self.state.lock().unwrap(); + state + .legacy_cancellation_requested + .insert(call_id.to_string()); + state + .active + .values() + .filter(|active| active.call_id == call_id) + .map(|active| active.handle.clone()) + .collect::>() + }; + for handle in handles { + self.session + .cancel_command(handle) + .await + .map_err(crate::ToolsError::from)?; + } + Ok(()) + } + + async fn cancel_execution( + &self, + ctx: &agen::tool::ToolExecutionContext, + ) -> Result<(), ToolError> { + let execution_id = ctx.execution_id(); + let handle = { + let mut state = self.state.lock().unwrap(); + state.cancellation_requested.insert(execution_id.clone()); + state + .active + .get(&execution_id) + .map(|active| active.handle.clone()) + }; + if let Some(handle) = handle { + self.session + .cancel_command(handle) + .await + .map_err(crate::ToolsError::from)?; + } + Ok(()) } } @@ -123,6 +267,7 @@ pub fn bash_tool(session: WorkdirSessionHandle, _output_dir: PathBuf) -> ToolDef .input_schema(serde_json::to_value(schema).expect("Bash schema serialization")); let tool: Arc = Arc::new(BashTool { session: session.clone(), + state: Arc::new(Mutex::new(BashExecutionState::default())), }); (meta, tool) }) diff --git a/crates/tools/src/grep.rs b/crates/tools/src/grep.rs index 598015b6..27c26ad0 100644 --- a/crates/tools/src/grep.rs +++ b/crates/tools/src/grep.rs @@ -129,7 +129,7 @@ pub fn grep_tool(session: WorkdirSessionHandle) -> ToolDefinition { Arc::new(move || { let schema = schemars::schema_for!(GrepParams); let meta = ToolMeta::new("Grep") - .description("Search Workdir file contents with a regex. Glob/Grep traversal executes inside the WorkdirSession provider. Results are bounded and Workdir-relative.") + .description("Search Workdir file contents with a regex. Content results group lines by file; `>` marks matching lines and unmarked lines are context. Glob/Grep traversal executes inside the WorkdirSession provider. Results are bounded and Workdir-relative.") .input_schema(serde_json::to_value(schema).expect("Grep schema serialization")); let tool: Arc = Arc::new(GrepTool { session: session.clone(), diff --git a/crates/tools/tests/integration.rs b/crates/tools/tests/integration.rs index f8761aa8..98e2dab6 100644 --- a/crates/tools/tests/integration.rs +++ b/crates/tools/tests/integration.rs @@ -7,7 +7,10 @@ use std::path::Path; use std::sync::Arc; -use agen::tool::{Tool, ToolDefinition, ToolMeta}; +use agen::tool::{ + Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolExecutionHandle, + ToolExecutionTerminal, ToolMeta, +}; use manifest::{Permission, Scope, ScopeConfig, ScopeRule}; use serde_json::json; use tempfile::TempDir; @@ -401,5 +404,84 @@ async fn bash_provider_output_does_not_expose_internal_paths() { assert_eq!(std::fs::read_dir(spill.path()).unwrap().count(), 0); } +#[tokio::test] +async fn bash_cancellation_returns_bounded_progress_as_terminal_output() { + let (dir, _spill, reg) = setup(); + let marker = dir.path().join("must-not-run-after-cancel"); + let command = format!( + "printf 'before\\n'; printf 'err-before\\n' >&2; sleep 1; touch {}; printf 'after\\n'", + marker.display() + ); + let input = serde_json::to_string(&json!({ "command": command })).unwrap(); + let context = ToolExecutionContext::new("call-heavy", "attempt-heavy", 0); + let bash = reg.get("Bash"); + let executing = bash.clone(); + let execution_context = context.clone(); + let execution = tokio::spawn(async move { executing.execute(&input, execution_context).await }); + + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + bash.cancel_execution(&context) + .await + .expect("signal exact execution cancellation"); + let error = tokio::time::timeout(std::time::Duration::from_secs(2), execution) + .await + .expect("cancelled Bash should terminate inside the Engine grace budget") + .expect("Bash task join"); + + let ToolError::Cancelled(output) = error.expect_err("cancelled command is non-success") else { + panic!("expected typed cancellation result"); + }; + let content = output.content.expect("bounded progress output"); + assert!( + content.contains("before"), + "missing pre-cancel stdout: {content}" + ); + assert!( + content.contains("err-before"), + "missing pre-cancel stderr: {content}" + ); + assert!( + !content.contains("after"), + "post-cancel output leaked: {content}" + ); + assert!(content.len() <= 16 * 1024, "output must remain bounded"); + + tokio::time::sleep(std::time::Duration::from_millis(1_100)).await; + assert!( + !marker.exists(), + "the cancelled command continued executing after terminal confirmation" + ); +} + +#[tokio::test] +async fn bash_force_close_cleanup_stops_command_and_keeps_session_reusable() { + let (dir, _spill, reg) = setup(); + let marker = dir.path().join("must-not-survive-force-close"); + let command = format!("sleep 1; touch {}", marker.display()); + let input = serde_json::to_string(&json!({ "command": command })).unwrap(); + let bash = reg.get("Bash"); + let context = ToolExecutionContext::new("call-force", "attempt-force", 0); + let (handle, terminal) = ToolExecutionHandle::start(bash.clone(), input, context); + + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + handle.force_close(); + assert!(matches!( + terminal.await, + ToolExecutionTerminal::OutcomeUnknown + )); + + tokio::time::sleep(std::time::Duration::from_millis(1_100)).await; + assert!( + !marker.exists(), + "CommandGuard cleanup allowed a force-closed command to continue" + ); + + let output = bash + .execute(r#"{"command":"printf 'reused'"}"#, Default::default()) + .await + .expect("workdir session remains reusable after cleanup"); + assert_eq!(output.content.as_deref(), Some("reused")); +} + // Sanity: unused Path import guard const _: fn() -> &'static Path = || Path::new("/"); diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index 2a8ee1bb..fd8b7bc9 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -1244,6 +1244,7 @@ impl App { id, summary, output, + disposition: _, is_error, } => { self.latest_llm_wait_event = None; diff --git a/crates/tui/src/console/mod.rs b/crates/tui/src/console/mod.rs index e7991a9c..0a6262ee 100644 --- a/crates/tui/src/console/mod.rs +++ b/crates/tui/src/console/mod.rs @@ -1016,7 +1016,7 @@ fn handle_key(app: &mut App, key: KeyEvent) -> Option { app.clear_queued_inputs(); Some(Method::Cancel) } - WorkerStatus::Idle => Some(Method::Shutdown), + WorkerStatus::Idle | WorkerStatus::Stopped => Some(Method::Shutdown), }), KeyCode::Char('d') if ctrl => { app.quit = true; diff --git a/crates/tui/src/dashboard/mod.rs b/crates/tui/src/dashboard/mod.rs index 7216a8c4..e56e272b 100644 --- a/crates/tui/src/dashboard/mod.rs +++ b/crates/tui/src/dashboard/mod.rs @@ -29,7 +29,6 @@ use ratatui::layout::{Constraint, Layout, Position, Rect}; use ratatui::style::{Color, Modifier, Style}; use ratatui::text::{Line, Span}; use ratatui::widgets::{Block, Borders, Clear, Paragraph, Widget, Wrap}; -use serde::Serialize; use session_store::FsStore; use session_store::FsWorkerStore; use ticket::config::{GitBranchName, TicketConfig, TicketOrchestrationConfig}; @@ -70,10 +69,6 @@ use render::{PanelListRow, row_hit_boxes}; const MAX_ENTRIES: usize = 50; const CLOSED_VISIBLE_ROWS: usize = 3; -const ORCHESTRATOR_IDLE_QUEUE_NOTICE_PROMPT: &str = "panel.orchestrator_idle_queue_notice"; -const ORCHESTRATOR_QUEUE_ATTENTION_MAX_TICKETS: usize = 6; -const ORCHESTRATOR_QUEUE_ATTENTION_MAX_TEXT_CHARS: usize = 120; -const ORCHESTRATOR_QUEUE_ATTENTION_MAX_MESSAGE_CHARS: usize = 2_400; const SOCKET_OP_TIMEOUT: Duration = Duration::from_secs(3); const DASHBOARD_POLL_INTERVAL: Duration = Duration::from_millis(1_500); const TERMINAL_EVENT_POLL_INTERVAL: Duration = Duration::from_millis(100); @@ -911,6 +906,7 @@ struct OrchestratorActiveWorkItem { #[derive(Debug, Clone, PartialEq, Eq)] struct OrchestratorQueuedWorkItem { id: String, + resource_key: Option, title: String, classification: OrchestratorQueuedClassification, waiting_reason: Option, @@ -975,22 +971,6 @@ impl OrchestratorQueueAttentionNoticeResult { } } -#[derive(Debug, Serialize)] -struct OrchestratorQueueTemplateContext { - workspace: String, - actionable_tickets: Vec, - waiting_tickets: Vec, - omitted_ticket_count: usize, -} - -#[derive(Debug, Serialize)] -struct OrchestratorQueueTemplateTicket { - id: String, - title: String, - classification: &'static str, - waiting_reason: Option, -} - #[derive(Debug, Clone, PartialEq, Eq)] struct PanelRowHitBox { rect: Rect, @@ -1326,7 +1306,16 @@ impl DashboardApp { if self.orchestrator_work_set.is_empty() { self.refresh_orchestrator_work_set(); } - let notice = orchestrator_queue_attention_notice(&self.panel, &self.orchestrator_work_set)?; + let notice = match orchestrator_queue_attention_notice(&self.orchestrator_work_set) { + Ok(Some(notice)) => notice, + Ok(None) => return None, + Err(error) => { + self.notice = Some(format!( + "Orchestrator queued-work attention not delivered: {error}" + )); + return None; + } + }; if self .orchestrator_queue_attention .as_ref() @@ -3661,6 +3650,7 @@ fn derive_orchestrator_work_set( }; Some(OrchestratorQueuedWorkItem { id: ticket.id.clone(), + resource_key: ticket.resource_key.clone(), title: ticket.title.clone(), classification, waiting_reason, @@ -3744,72 +3734,46 @@ fn orchestrator_work_set_fingerprint( } fn orchestrator_queue_attention_notice( - panel: &WorkspacePanelViewModel, work_set: &OrchestratorWorkSet, -) -> Option { +) -> Result, &'static str> { if work_set.has_active_inprogress() { - return None; + return Ok(None); } let actionable = work_set.actionable_queued(); if actionable.is_empty() { - return None; + return Ok(None); } let waiting = work_set .queued .iter() - .filter(|item| item.waiting_reason.is_some()) - .collect::>(); - let ticket_count = actionable.len() + waiting.len(); - let actionable_tickets = actionable - .iter() - .take(ORCHESTRATOR_QUEUE_ATTENTION_MAX_TICKETS) - .map(|item| orchestrator_queue_template_ticket(item)) - .collect::>(); - let remaining_capacity = - ORCHESTRATOR_QUEUE_ATTENTION_MAX_TICKETS.saturating_sub(actionable_tickets.len()); - let waiting_tickets = waiting - .iter() - .take(remaining_capacity) - .map(|item| orchestrator_queue_template_ticket(item)) - .collect::>(); - let rendered = - render_orchestrator_queue_attention_template(&OrchestratorQueueTemplateContext { - workspace: bounded_progress_text( - &panel.header.workspace_label, - ORCHESTRATOR_QUEUE_ATTENTION_MAX_TEXT_CHARS, - ), - actionable_tickets, - waiting_tickets, - omitted_ticket_count: ticket_count - .saturating_sub(ORCHESTRATOR_QUEUE_ATTENTION_MAX_TICKETS), + .filter(|item| item.waiting_reason.is_some()); + let tickets = actionable + .into_iter() + .chain(waiting) + .map(|item| { + let resource_key = item + .resource_key + .clone() + .ok_or("queued Ticket is missing its required resource key")?; + worker::OrchestratorQueueAttentionTicket::new(resource_key, item.title.clone()) + .map_err(|_| "queued Ticket has an invalid resource key") }) - .ok()?; - let message = bounded_progress_text(&rendered, ORCHESTRATOR_QUEUE_ATTENTION_MAX_MESSAGE_CHARS); + .collect::, _>>()?; + let context = worker::OrchestratorQueueAttentionContext::new(tickets); + let message = render_orchestrator_queue_attention_template(&context) + .map_err(|_| "queued-work attention prompt rendering failed")?; let fingerprint = format!("idle-queue:{}", work_set.fingerprint); - Some(OrchestratorQueueAttentionNotice { + Ok(Some(OrchestratorQueueAttentionNotice { message, fingerprint, - }) -} - -fn orchestrator_queue_template_ticket( - item: &&OrchestratorQueuedWorkItem, -) -> OrchestratorQueueTemplateTicket { - OrchestratorQueueTemplateTicket { - id: bounded_progress_text(&item.id, ORCHESTRATOR_QUEUE_ATTENTION_MAX_TEXT_CHARS), - title: bounded_progress_text(&item.title, ORCHESTRATOR_QUEUE_ATTENTION_MAX_TEXT_CHARS), - classification: item.classification.as_str(), - waiting_reason: item.waiting_reason.as_ref().map(|reason| { - bounded_progress_text(reason, ORCHESTRATOR_QUEUE_ATTENTION_MAX_TEXT_CHARS) - }), - } + })) } fn render_orchestrator_queue_attention_template( - context: &OrchestratorQueueTemplateContext, + context: &worker::OrchestratorQueueAttentionContext, ) -> Result { worker::PromptCatalog::builtins_only()? - .render_serializable(ORCHESTRATOR_IDLE_QUEUE_NOTICE_PROMPT, context) + .orchestrator_queue_attention(worker::OrchestratorQueueAttentionPrompt::Tui, context) } fn orchestrator_work_set_detail( @@ -5236,6 +5200,7 @@ fn row_status_label(entry: &WorkerListEntry) -> (&'static str, Style) { .fg(Color::Cyan) .add_modifier(Modifier::BOLD), ), + Some(WorkerStatus::Stopped) => ("live stopped", Style::default().fg(Color::DarkGray)), None => ("live", Style::default().fg(Color::DarkGray)), }; } diff --git a/crates/tui/src/dashboard/tests.rs b/crates/tui/src/dashboard/tests.rs index 62d72915..32be9393 100644 --- a/crates/tui/src/dashboard/tests.rs +++ b/crates/tui/src/dashboard/tests.rs @@ -2972,7 +2972,7 @@ fn dashboard_empty_enter_on_non_openable_row_reports_open_diagnostic() { } #[test] -fn idle_orchestrator_gets_bounded_attention_for_new_queued_work() { +fn idle_orchestrator_gets_sanitized_attention_for_new_queued_work() { let mut app = ticket_enabled_app(vec![live_info("test-orchestrator", WorkerStatus::Idle)]); app.panel.rows = vec![panel_test_ticket_row( "00001QUEUE", @@ -2992,11 +2992,87 @@ fn idle_orchestrator_gets_bounded_attention_for_new_queued_work() { request .notice .message - .starts_with("Workspace Dashboard observed") + .starts_with("Queued Tickets require attention:") ); - assert!(request.notice.message.contains("00001QUEUE")); - assert!(request.notice.message.contains("new_queued")); - assert!(request.notice.message.contains("queued -> inprogress")); + assert!(request.notice.message.contains("- T-1 — Queued work")); + assert!( + request + .notice + .message + .contains("Reread the current Ticket state before acting") + ); + assert!( + !request + .notice + .message + .contains(&app.panel.header.workspace_label) + ); + for hidden in [ + "00001QUEUE", + "Workspace:", + "workspace_id", + "new_queued", + "bounded", + "queued -> inprogress", + ] { + assert!(!request.notice.message.contains(hidden), "leaked {hidden}"); + } +} + +#[test] +fn queued_attention_missing_resource_key_fails_closed_with_panel_notice() { + let mut app = ticket_enabled_app(vec![live_info("test-orchestrator", WorkerStatus::Idle)]); + let mut row = panel_test_ticket_row( + "00001QUEUE", + "Queued work", + ActionPriority::Background, + NextUserAction::Wait, + "queued", + ); + row.ticket.as_mut().unwrap().resource_key = None; + app.panel.rows = vec![row]; + app.refresh_orchestrator_work_set(); + + assert!(app.prepare_orchestrator_queue_attention_notice().is_none()); + assert_eq!( + app.notice.as_deref(), + Some( + "Orchestrator queued-work attention not delivered: queued Ticket is missing its required resource key" + ) + ); +} + +#[test] +fn queued_attention_truncates_only_when_tickets_are_omitted() { + let mut app = ticket_enabled_app(vec![live_info("test-orchestrator", WorkerStatus::Idle)]); + app.panel.rows = (1..=worker::OrchestratorQueueAttentionContext::MAX_TICKETS + 1) + .map(|index| { + let mut row = panel_test_ticket_row( + &format!("opaque-{index}"), + &format!("Queued work {index}"), + ActionPriority::Background, + NextUserAction::Wait, + "queued", + ); + row.ticket.as_mut().unwrap().resource_key = Some(format!("T-{index}")); + row + }) + .collect(); + app.refresh_orchestrator_work_set(); + + let request = app + .prepare_orchestrator_queue_attention_notice() + .expect("bounded queued-work attention"); + + assert!(request.notice.message.contains("- T-20 — Queued work 20")); + assert!(!request.notice.message.contains("T-21")); + assert!( + request + .notice + .message + .contains("were omitted from this notice: 1") + ); + assert!(!request.notice.message.contains("opaque-")); } #[test] @@ -3086,7 +3162,9 @@ fn planned_queued_prompts_when_active_work_clears() { .prepare_orchestrator_queue_attention_notice() .expect("planned queued work should prompt after active work clears"); - assert!(request.notice.message.contains("planned_queued")); + assert!(request.notice.message.contains("- T-1 — Queued work")); + assert!(!request.notice.message.contains("planned_queued")); + assert!(!request.notice.message.contains("00001QUEUE")); assert!( !request .notice @@ -3141,8 +3219,9 @@ fn rediscovered_queued_work_is_actionable_when_session_work_set_is_empty() { .prepare_orchestrator_queue_attention_notice() .expect("queued ticket state should be rediscovered safely"); - assert!(request.notice.message.contains("new_queued")); - assert!(request.notice.message.contains("00001QUEUE")); + assert!(request.notice.message.contains("- T-1 — Queued work")); + assert!(!request.notice.message.contains("new_queued")); + assert!(!request.notice.message.contains("00001QUEUE")); } #[test] diff --git a/crates/tui/src/workspace_panel.rs b/crates/tui/src/workspace_panel.rs index fa2dcc40..8b3223fa 100644 --- a/crates/tui/src/workspace_panel.rs +++ b/crates/tui/src/workspace_panel.rs @@ -1530,6 +1530,7 @@ fn worker_status_label(entry: &WorkerListEntry) -> &'static str { Some(WorkerStatus::Idle) => "live idle", Some(WorkerStatus::Running) => "live running", Some(WorkerStatus::Paused) => "live paused", + Some(WorkerStatus::Stopped) => "live stopped", None => "live", }; } diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 5bc2ad5e..03f3b5f4 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -2742,6 +2742,7 @@ impl RuntimeState { protocol::WorkerStatus::Running => Some(WorkerStatus::Running), protocol::WorkerStatus::Idle => Some(WorkerStatus::Idle), protocol::WorkerStatus::Paused => Some(WorkerStatus::Paused), + protocol::WorkerStatus::Stopped => Some(WorkerStatus::Stopped), }, protocol::Event::RunEnd { result } => match result { protocol::RunResult::Finished | protocol::RunResult::RolledBack => { @@ -3104,7 +3105,7 @@ mod tests { &mut activity, &internal_worker_status_event( internal_worker_ref("child-b", None), - protocol::WorkerStatus::Idle, + protocol::WorkerStatus::Stopped, ), )); } diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index fe245c6f..d112d7df 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -38,9 +38,7 @@ use crate::working_directory::{ }; use async_trait::async_trait; use protocol::{Event, Method, Segment, WorkerStatus}; -use session_store::{ - CombinedStore, LogEntry, WorkerAggregateStore, WorkerSessionStore, collect_state, -}; +use session_store::{CombinedStore, LogEntry, WorkerAggregateStore, WorkerSessionStore}; #[cfg(test)] use session_store::{FsStore, FsWorkerStore}; use tokio::runtime::Runtime; @@ -68,8 +66,10 @@ const RUNTIME_TASK_TIMEOUT: Duration = Duration::from_secs(10); const USER_INPUT_COMMIT_TIMEOUT: Duration = Duration::from_secs(9); fn user_input_has_submission(entry: &LogEntry, submission_id: &str) -> bool { - let LogEntry::UserInput { extensions, .. } = entry else { - return false; + let extensions = match entry { + LogEntry::UserInput { extensions, .. } + | LogEntry::AnnotatedUserInput { extensions, .. } => extensions, + _ => return false, }; extensions.iter().any(|extension| { extension.domain == WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN @@ -212,11 +212,11 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider { return Err(WorkerObservationError::NotFound); } let entries = sink.subscribe_with_snapshot().0; - let state = collect_state(&entries); - Ok(WorkerSessionCapture { - segment_id: format!("runtime:{runtime_id}:worker:{worker_id}"), - items: state.history, - }) + WorkerSessionCapture::from_log_entries( + format!("runtime:{runtime_id}:worker:{worker_id}"), + &entries, + ) + .map_err(WorkerObservationError::Unavailable) } } @@ -1546,7 +1546,9 @@ fn accepted_notify_run_state(status: WorkerStatus, auto_run: bool) -> WorkerExec match status { WorkerStatus::Running => WorkerExecutionRunState::Busy, WorkerStatus::Idle if auto_run => WorkerExecutionRunState::Busy, - WorkerStatus::Idle | WorkerStatus::Paused => WorkerExecutionRunState::Idle, + WorkerStatus::Idle | WorkerStatus::Paused | WorkerStatus::Stopped => { + WorkerExecutionRunState::Idle + } } } @@ -2505,7 +2507,9 @@ mod tests { let scope = Scope::writable(&scope_root).map_err(|err| err.to_string())?; let worker = Worker::new( manifest, - Engine::new(self.client.clone()), + Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated( + self.client.clone(), + ), store, workspace_context, filesystem_authority, @@ -3241,14 +3245,17 @@ mod tests { matches!( entry, LogEntry::UserInput { segments, .. } + | LogEntry::AnnotatedUserInput { segments, .. } if segments == &vec![Segment::text("start the ticket")] ) })); let submission_id = entries .iter() .find_map(|entry| { - let LogEntry::UserInput { extensions, .. } = entry else { - return None; + let extensions = match entry { + LogEntry::UserInput { extensions, .. } + | LogEntry::AnnotatedUserInput { extensions, .. } => extensions, + _ => return None, }; extensions .iter() diff --git a/crates/worker/examples/worker_cli.rs b/crates/worker/examples/worker_cli.rs index d7e26f0c..b14e45df 100644 --- a/crates/worker/examples/worker_cli.rs +++ b/crates/worker/examples/worker_cli.rs @@ -66,11 +66,12 @@ async fn main() -> Result<(), Box> { WorkerRunResult::Finished => println!("(finished)"), WorkerRunResult::Paused => println!("(paused)"), WorkerRunResult::LimitReached => println!("(turn limit reached)"), + WorkerRunResult::Interrupted { message, .. } => println!("(interrupted: {message})"), WorkerRunResult::RolledBack => println!("(empty turn rolled back)"), } // 5. Extract the assistant's reply from history - let history = worker.engine().history(); + let history = worker.history(); if let Some(text) = history .iter() .rev() diff --git a/crates/worker/src/compact/prune.rs b/crates/worker/src/compact/prune.rs index 0bba4395..8288bbf0 100644 --- a/crates/worker/src/compact/prune.rs +++ b/crates/worker/src/compact/prune.rs @@ -22,7 +22,7 @@ use crate::compact::token_counter::{ EstimateSource, savings_for_prune_impl, token_estimates_for_prune_impl, }; -impl Worker { +impl Worker { /// Enable prune projection on the underlying Engine. /// /// Registers the config and token/savings-estimator closures on the Engine. diff --git a/crates/worker/src/compact/token_counter.rs b/crates/worker/src/compact/token_counter.rs index 19cf52de..dff10194 100644 --- a/crates/worker/src/compact/token_counter.rs +++ b/crates/worker/src/compact/token_counter.rs @@ -242,13 +242,13 @@ pub(crate) fn savings_for_prune_impl( // ── Worker に生やす公開 API ─────────────────────────────────────────────── -impl Worker { +impl Worker { /// 現在の history 全体の推定トークン数。 /// /// 最後の measurement と、その後に追加された未測定分の byte/4 外挿。 pub fn total_tokens(&self) -> TokenEstimate { let usage = self.usage_history(); - agen::token_counter::total_tokens(self.history(), &usage) + agen::token_counter::total_tokens(&self.history(), &usage) } /// 任意の history index 時点でのプロンプト全長推定。 @@ -259,7 +259,7 @@ impl Worker { /// pointer 以降に増えたプロンプト長を測るのに使う。 pub fn total_tokens_at(&self, history_len: usize) -> TokenEstimate { let usage = self.usage_history(); - agen::token_counter::total_tokens_at(self.history(), &usage, history_len) + agen::token_counter::total_tokens_at(&self.history(), &usage, history_len) } /// 末尾から `retained` トークン以上を残すための分割位置。 @@ -267,7 +267,7 @@ impl Worker { /// `history[..cut.index]` が要約/破棄される側、`history[cut.index..]` が残る側。 pub fn split_for_retained(&self, retained: u64) -> SplitPoint { let usage = self.usage_history(); - split_for_retained_impl(self.history(), &usage, retained) + split_for_retained_impl(&self.history(), &usage, retained) } } diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 8e538bb3..ff943c60 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -485,6 +485,7 @@ impl WorkerController { // into the controller task so the in-flight turn can be reached // via these handles while worker itself is borrowed by drive_turn. let cancel_tx = worker.engine_mut().cancel_sender(); + let pause_tx = worker.engine_mut().pause_sender(); let notify_buffer = worker.notify_buffer_handle(); tokio::spawn(controller_loop( @@ -494,6 +495,7 @@ impl WorkerController { shared_state, runtime_dir, cancel_tx, + pause_tx, notify_buffer, self_parent_socket, spawner_name, @@ -763,6 +765,19 @@ pub(crate) fn wire_event_bridges_on_engine( id: result.tool_use_id.clone(), summary: result.summary.clone(), output: result.content.clone(), + disposition: Some(match result.disposition { + agen::ToolResultDisposition::Success => protocol::ToolResultDisposition::Success, + agen::ToolResultDisposition::Error => protocol::ToolResultDisposition::Error, + agen::ToolResultDisposition::Interrupted => { + protocol::ToolResultDisposition::Interrupted + } + agen::ToolResultDisposition::Cancelled => { + protocol::ToolResultDisposition::Cancelled + } + agen::ToolResultDisposition::OutcomeUnknown => { + protocol::ToolResultDisposition::OutcomeUnknown + } + }), is_error: result.is_error, }); }); @@ -1115,6 +1130,7 @@ async fn controller_loop( shared_state: Arc, runtime_dir: Arc, cancel_tx: mpsc::Sender<()>, + pause_tx: mpsc::Sender<()>, notify_buffer: NotifyBuffer, self_parent_socket: Option, spawner_name: String, @@ -1161,22 +1177,35 @@ async fn controller_loop( // clear at run start prevents stale partial output left by an older // interrupted/error turn from being carried into the next snapshot. worker.clear_in_flight_events(); - set_controller_status( - &shared_state, - &runtime_dir, - &event_tx, - WorkerStatus::Running, - ) - .await; let parent_originated = run.is_parent_originated(); + let user_input_run = matches!(&run, PendingRun::Run(_) | PendingRun::RunTracked { .. }); + if !user_input_run { + set_controller_status( + &shared_state, + &runtime_dir, + &event_tx, + WorkerStatus::Running, + ) + .await; + } let (mut new_status, shutdown) = match run { PendingRun::Run(input) => { + let (input_commit_tx, input_commit_rx) = oneshot::channel(); drive_turn( - worker.run(input), + worker.run_with_input_extensions_and_commit_hook( + input, + Vec::new(), + move || { + let _ = input_commit_tx.send(()); + }, + ), &mut method_rx, &event_tx, &cancel_tx, + &pause_tx, &shared_state, + &runtime_dir, + Some(input_commit_rx), ¬ify_buffer, self_parent_socket.as_ref(), &spawner_name, @@ -1186,12 +1215,22 @@ async fn controller_loop( .await } PendingRun::RunTracked { input, extension } => { + let (input_commit_tx, input_commit_rx) = oneshot::channel(); drive_turn( - worker.run_with_input_extensions(input, vec![extension]), + worker.run_with_input_extensions_and_commit_hook( + input, + vec![extension], + move || { + let _ = input_commit_tx.send(()); + }, + ), &mut method_rx, &event_tx, &cancel_tx, + &pause_tx, &shared_state, + &runtime_dir, + Some(input_commit_rx), ¬ify_buffer, self_parent_socket.as_ref(), &spawner_name, @@ -1206,7 +1245,10 @@ async fn controller_loop( &mut method_rx, &event_tx, &cancel_tx, + &pause_tx, &shared_state, + &runtime_dir, + None, ¬ify_buffer, self_parent_socket.as_ref(), &spawner_name, @@ -1221,7 +1263,10 @@ async fn controller_loop( &mut method_rx, &event_tx, &cancel_tx, + &pause_tx, &shared_state, + &runtime_dir, + None, ¬ify_buffer, self_parent_socket.as_ref(), &spawner_name, @@ -1346,7 +1391,7 @@ async fn controller_loop( }); } }, - WorkerStatus::Idle => { + WorkerStatus::Idle | WorkerStatus::Stopped => { let _ = event_tx.send(Event::Error { code: ErrorCode::NotRunning, message: "Worker is not running".into(), @@ -1387,7 +1432,7 @@ async fn controller_loop( .into(), }); } - WorkerStatus::Running => { + WorkerStatus::Running | WorkerStatus::Stopped => { let _ = event_tx.send(Event::Error { code: ErrorCode::AlreadyRunning, message: @@ -1401,7 +1446,7 @@ async fn controller_loop( WorkerStatus::Idle | WorkerStatus::Paused => { emit_rewind_targets(&worker, &event_tx) } - WorkerStatus::Running => { + WorkerStatus::Running | WorkerStatus::Stopped => { let _ = event_tx.send(Event::Error { code: ErrorCode::AlreadyRunning, message: "Worker is already executing a turn; rewind can only run while idle or paused" @@ -1430,7 +1475,7 @@ async fn controller_loop( .into(), }); } - WorkerStatus::Running => { + WorkerStatus::Running | WorkerStatus::Stopped => { let _ = event_tx.send(Event::Error { code: ErrorCode::AlreadyRunning, message: "Worker is already executing a turn; rewind can only run while idle or paused" @@ -1618,7 +1663,10 @@ async fn drive_turn( method_rx: &mut mpsc::Receiver, event_tx: &broadcast::Sender, cancel_tx: &mpsc::Sender<()>, + pause_tx: &mpsc::Sender<()>, shared_state: &Arc, + runtime_dir: &RuntimeDir, + mut input_commit_rx: Option>, notify_buffer: &NotifyBuffer, parent_socket: Option<&PathBuf>, self_name: &str, @@ -1634,14 +1682,58 @@ where loop { tokio::select! { + // If input commit and provider completion become ready together, expose + // Running only after processing the commit fence. This makes the + // Running snapshot contract deterministic even for immediate clients. + biased; + committed = async { + input_commit_rx + .as_mut() + .expect("input commit receiver guarded by select condition") + .await + }, if input_commit_rx.is_some() => { + input_commit_rx = None; + if committed.is_ok() { + set_controller_status( + shared_state, + runtime_dir, + event_tx, + WorkerStatus::Running, + ) + .await; + } + } result = &mut worker_future => { return match result { Ok(r) => { let (status, run_result) = match r { + WorkerRunResult::Finished if pause_requested => { + (WorkerStatus::Paused, RunResult::Paused) + } WorkerRunResult::Finished => (WorkerStatus::Idle, RunResult::Finished), WorkerRunResult::Paused => (WorkerStatus::Paused, RunResult::Paused), WorkerRunResult::LimitReached => (WorkerStatus::Idle, RunResult::LimitReached), WorkerRunResult::RolledBack => (WorkerStatus::Idle, RunResult::RolledBack), + WorkerRunResult::Interrupted { .. } if pause_requested => { + let _ = event_tx.send(Event::RunEnd { result: RunResult::Paused }); + return (WorkerStatus::Paused, shutdown_requested); + } + WorkerRunResult::Interrupted { code, message } => { + let _ = event_tx.send(Event::Error { + code, + message: message.clone(), + }); + if parent_originated { + crate::ipc::event::fire_and_forget( + parent_socket.cloned(), + protocol::WorkerEvent::Errored { + worker_name: self_name.to_string(), + message, + }, + ); + } + return (WorkerStatus::Idle, shutdown_requested); + } }; let _ = event_tx.send(Event::RunEnd { result: run_result }); if parent_originated && matches!(run_result, RunResult::Finished) { @@ -1690,7 +1782,7 @@ where } Some(Method::Pause) => { pause_requested = true; - let _ = cancel_tx.try_send(()); + let _ = pause_tx.try_send(()); } Some(Method::Shutdown) => { shutdown_requested = true; @@ -1752,7 +1844,7 @@ where fn emit_rewind_targets(worker: &Worker, event_tx: &broadcast::Sender) where - C: LlmClient, + C: LlmClient + 'static, St: Store, { match worker.list_rewind_targets() { @@ -1778,7 +1870,7 @@ fn apply_rewind( expected_head_entries: usize, ) -> bool where - C: LlmClient, + C: LlmClient + 'static, St: Store, { match worker.rewind_to(target, expected_head_entries) { @@ -1826,7 +1918,7 @@ fn model_supports_image_attachments(model: &manifest::ModelManifest) -> bool { fn build_greeting(worker: &Worker) -> protocol::Greeting where - C: LlmClient, + C: LlmClient + 'static, St: Store, { let manifest = worker.manifest(); @@ -1942,11 +2034,13 @@ mod tests { event_tx: broadcast::Sender, cancel_tx: mpsc::Sender<()>, _cancel_rx: mpsc::Receiver<()>, + pause_tx: mpsc::Sender<()>, + _pause_rx: mpsc::Receiver<()>, shared_state: Arc, notify_buffer: NotifyBuffer, spawned_registry: Arc, parent_socket_path: PathBuf, - _runtime_dir: Arc, + runtime_dir: Arc, _temp: TempDir, } @@ -1960,6 +2054,7 @@ mod tests { let (method_tx, method_rx) = mpsc::channel::(16); let (event_tx, _) = broadcast::channel::(16); let (cancel_tx, cancel_rx) = mpsc::channel::<()>(1); + let (pause_tx, pause_rx) = mpsc::channel::<()>(1); let shared_state = Arc::new(WorkerSharedState::new( "child-worker".to_string(), session_store::new_segment_id(), @@ -1985,11 +2080,13 @@ mod tests { event_tx, cancel_tx, _cancel_rx: cancel_rx, + pause_tx, + _pause_rx: pause_rx, shared_state, notify_buffer, spawned_registry, parent_socket_path, - _runtime_dir: runtime_dir, + runtime_dir, _temp: temp, } } @@ -2042,7 +2139,10 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, + &env.runtime_dir, + None, &env.notify_buffer, Some(&env.parent_socket_path), "child-worker", @@ -2063,6 +2163,44 @@ mod tests { } } + #[tokio::test] + async fn pause_waits_for_run_boundary_and_uses_safe_pause_channel() { + let mut env = make_env().await; + let method_tx = env._method_tx.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(10)).await; + method_tx.send(Method::Pause).await.expect("send pause"); + }); + + let worker_future = async { + tokio::time::sleep(Duration::from_millis(100)).await; + Ok::<_, WorkerError>(WorkerRunResult::Finished) + }; + let started_at = std::time::Instant::now(); + let (status, shutdown) = drive_turn( + worker_future, + &mut env.method_rx, + &env.event_tx, + &env.cancel_tx, + &env.pause_tx, + &env.shared_state, + &env.runtime_dir, + None, + &env.notify_buffer, + None, + "child-worker", + &env.spawned_registry, + true, + ) + .await; + + assert_eq!(status, WorkerStatus::Paused); + assert!(!shutdown); + assert!(started_at.elapsed() >= Duration::from_millis(100)); + assert!(env._pause_rx.try_recv().is_ok()); + assert!(env._cancel_rx.try_recv().is_err()); + } + #[tokio::test] async fn non_parent_originated_finished_stays_silent() { let mut env = make_env().await; @@ -2074,7 +2212,10 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, + &env.runtime_dir, + None, &env.notify_buffer, Some(&env.parent_socket_path), "child-worker", @@ -2109,7 +2250,10 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, + &env.runtime_dir, + None, &env.notify_buffer, Some(&env.parent_socket_path), "child-worker", @@ -2150,7 +2294,10 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, + &env.runtime_dir, + None, &env.notify_buffer, Some(&env.parent_socket_path), "child-worker", @@ -2189,7 +2336,10 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, + &env.runtime_dir, + None, &env.notify_buffer, Some(&env.parent_socket_path), "parent", @@ -2225,7 +2375,10 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, + &env.runtime_dir, + None, &env.notify_buffer, Some(&env.parent_socket_path), "parent", @@ -2259,7 +2412,10 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, + &env.runtime_dir, + None, &env.notify_buffer, Some(&env.parent_socket_path), "parent", @@ -2292,7 +2448,10 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, + &env.runtime_dir, + None, &env.notify_buffer, Some(&env.parent_socket_path), "child-worker", diff --git a/crates/worker/src/feature.rs b/crates/worker/src/feature.rs index 54845c25..f0231d15 100644 --- a/crates/worker/src/feature.rs +++ b/crates/worker/src/feature.rs @@ -1795,9 +1795,9 @@ impl FeatureRegistryBuilder { } /// Install modules into the existing Engine tool path and hook builder. - pub(crate) fn install_into_engine( + pub(crate) fn install_into_engine( self, - worker: &mut Engine, + worker: &mut Engine, hook_builder: &mut HookRegistryBuilder, ) -> FeatureRegistryInstallReport { let mut pending_tools = Vec::new(); diff --git a/crates/worker/src/feature/builtin.rs b/crates/worker/src/feature/builtin.rs index b9c5994f..010fcf96 100644 --- a/crates/worker/src/feature/builtin.rs +++ b/crates/worker/src/feature/builtin.rs @@ -12,6 +12,7 @@ pub mod memory_extract; pub mod merge_request; pub mod objective; pub mod orchestration; +mod resource_projection; pub mod session_explore; pub mod task; pub mod ticket; diff --git a/crates/worker/src/feature/builtin/memory_extract.rs b/crates/worker/src/feature/builtin/memory_extract.rs index 61400fe7..165efe6a 100644 --- a/crates/worker/src/feature/builtin/memory_extract.rs +++ b/crates/worker/src/feature/builtin/memory_extract.rs @@ -6,7 +6,9 @@ use memory::backend::{ MemoryBackendOperation, MemoryBackendOperationResult, MemoryStageCandidateOperation, }; use memory::extract::{CandidateKind, ExtractedCandidate, StagingEvidence}; -use memory::schema::{EvidenceKind, SourceEvidenceRef, SourceRef}; +use memory::schema::{ + EvidenceKind, EvidenceOrigin, EvidenceOriginKind, SourceEvidenceRef, SourceRef, +}; use schemars::JsonSchema; use serde::Deserialize; @@ -174,17 +176,29 @@ impl Tool for StageMemoryCandidateTool { "StageMemoryCandidate requires at least one entry_ref".to_string(), )); } - let mut evidence = Vec::with_capacity(params.entry_refs.len()); - let mut source_refs = Vec::with_capacity(params.entry_refs.len()); + let mut entries = Vec::with_capacity(params.entry_refs.len()); for entry_ref in ¶ms.entry_refs { - let projection = self.state.view.evidence_for(entry_ref).ok_or_else(|| { + entries.push(self.state.view.evidence_for(entry_ref).ok_or_else(|| { ToolError::InvalidArgument(format!( "unknown SessionEntryRef {entry_ref:?} for this extraction capture" )) - })?; - evidence.push(staging_evidence(&projection)); - source_refs.push(source_evidence_ref(&projection)); + })?); } + if matches!(params.kind, CandidateKind::Preference) + && entries.iter().any(|entry| { + !matches!( + entry.origin, + crate::WorkerHistoryProvenance::HumanInput { .. } + ) + }) + { + return Err(ToolError::InvalidArgument( + "preference candidates require exclusively HumanInput evidence; model, Worker, Flow, backend, derived, and legacy-unknown origins are not preference authority" + .to_string(), + )); + } + let evidence = entries.iter().map(staging_evidence).collect(); + let source_refs = entries.iter().map(source_evidence_ref).collect(); let candidate = ExtractedCandidate { kind: params.kind, claim: params.claim, @@ -310,11 +324,65 @@ fn evidence_kind(entry: &SessionEntryEvidence) -> EvidenceKind { } } +fn evidence_origin(origin: &crate::WorkerHistoryProvenance) -> EvidenceOrigin { + use crate::WorkerHistoryProvenance as Origin; + let mut evidence = EvidenceOrigin { + kind: EvidenceOriginKind::LegacyUnknown, + account_id: None, + workspace_id: None, + runtime_id: None, + worker_id: None, + flow_selector: None, + flow_definition_id: None, + flow_definition_revision: None, + }; + match origin { + Origin::HumanInput { account_id } => { + evidence.kind = EvidenceOriginKind::HumanInput; + evidence.account_id = Some(account_id.clone()); + } + Origin::WorkerInput { actor } => { + evidence.kind = EvidenceOriginKind::WorkerInput; + evidence.workspace_id = actor.workspace_id.clone(); + evidence.runtime_id = actor.runtime_id.clone(); + evidence.worker_id = Some(actor.worker_id.clone()); + } + Origin::FlowInstruction { + selector, + definition_id, + definition_revision, + .. + } => { + evidence.kind = EvidenceOriginKind::FlowInstruction; + evidence.flow_selector = Some(selector.clone()); + evidence.flow_definition_id = Some(definition_id.clone()); + evidence.flow_definition_revision = Some(*definition_revision); + } + Origin::BackendInstruction { .. } => evidence.kind = EvidenceOriginKind::BackendInstruction, + Origin::ModelOutput { worker } => { + evidence.kind = EvidenceOriginKind::ModelOutput; + evidence.workspace_id = worker.workspace_id.clone(); + evidence.runtime_id = worker.runtime_id.clone(); + evidence.worker_id = Some(worker.worker_id.clone()); + } + Origin::ToolOutput { worker } => { + evidence.kind = EvidenceOriginKind::ToolOutput; + evidence.workspace_id = worker.workspace_id.clone(); + evidence.runtime_id = worker.runtime_id.clone(); + evidence.worker_id = Some(worker.worker_id.clone()); + } + Origin::DerivedSummary => evidence.kind = EvidenceOriginKind::DerivedSummary, + Origin::LegacyUnknown => evidence.kind = EvidenceOriginKind::LegacyUnknown, + } + evidence +} + fn staging_evidence(entry: &SessionEntryEvidence) -> StagingEvidence { StagingEvidence { id: entry.entry_ref.to_string(), kind: evidence_kind(entry), entry_range: Some(entry.entry_range), + origin: Some(evidence_origin(&entry.origin)), excerpt: Some(entry.excerpt.clone()), summary: Some(entry.summary.clone()), } @@ -325,6 +393,7 @@ fn source_evidence_ref(entry: &SessionEntryEvidence) -> SourceEvidenceRef { segment_id: Some(entry.segment_id.clone()), entry_range: Some(entry.entry_range), evidence_id: Some(entry.entry_ref.to_string()), + origin: Some(evidence_origin(&entry.origin)), evidence_kind: Some(evidence_kind(entry)), label: Some(entry.label.clone()), summary: Some(entry.summary.clone()), @@ -432,6 +501,15 @@ mod tests { assert!(input.contains("StageMemoryCandidate.entry_refs")); } + #[test] + fn human_origin_projects_account_authority_into_evidence() { + let origin = evidence_origin(&crate::WorkerHistoryProvenance::HumanInput { + account_id: "account-1".into(), + }); + assert_eq!(origin.kind, EvidenceOriginKind::HumanInput); + assert_eq!(origin.account_id.as_deref(), Some("account-1")); + } + #[test] fn backend_input_failures_remain_invalid_argument_tool_errors() { let backend = map_memory_stage_error(WorkspaceMemoryBackendError::Backend( @@ -445,6 +523,19 @@ mod tests { assert!(matches!(http, ToolError::InvalidArgument(_))); } + #[tokio::test] + async fn preference_rejects_legacy_unknown_before_backend_mutation() { + let tool = StageMemoryCandidateTool { state: state() }; + let error = tool + .execute( + r#"{"kind":"preference","claim":"claim","why_useful":"useful","entry_refs":["E00000000"]}"#, + agen::tool::ToolExecutionContext::direct(), + ) + .await + .unwrap_err(); + assert!(format!("{error:?}").contains("exclusively HumanInput evidence")); + } + #[tokio::test] async fn stage_rejects_entry_ref_outside_capture_before_backend_mutation() { let tool = StageMemoryCandidateTool { state: state() }; diff --git a/crates/worker/src/feature/builtin/objective.rs b/crates/worker/src/feature/builtin/objective.rs index 460c929a..2c91cefd 100644 --- a/crates/worker/src/feature/builtin/objective.rs +++ b/crates/worker/src/feature/builtin/objective.rs @@ -14,6 +14,8 @@ use serde_json::json; use crate::worker::{WorkspaceClient, WorkspaceRequest, WorkspaceRequestMethod}; +use super::resource_projection::{project_objective_detail, project_objective_query}; + #[derive(Clone, Debug)] pub struct WorkspaceHttpObjectiveBackend { client: Arc, @@ -37,6 +39,7 @@ impl WorkspaceHttpObjectiveBackend { ) .await .map_err(backend_error)?; + let response = project_objective_query(response).map_err(ToolError::ExecutionFailed)?; Ok(ToolOutput { summary: "Queried Objectives".to_string(), content: Some(serde_json::to_string_pretty(&response).map_err(decode_error)?), @@ -58,8 +61,10 @@ impl WorkspaceHttpObjectiveBackend { ) .await .map_err(backend_error)?; + let response = project_objective_detail(response).map_err(ToolError::ExecutionFailed)?; + let objective_ref = response.objective_ref().to_string(); Ok(ToolOutput { - summary: format!("Read objective {id}"), + summary: format!("Read objective {objective_ref}"), content: Some(serde_json::to_string_pretty(&response).map_err(decode_error)?), attachments: Vec::new(), }) @@ -84,7 +89,7 @@ impl WorkspaceHttpObjectiveBackend { .await .map_err(backend_error)?; Ok(objective_output( - format!("Created objective {}", response.id), + format!("Created objective {}", &response.resource_key), response, )?) } @@ -112,7 +117,7 @@ impl WorkspaceHttpObjectiveBackend { .await .map_err(backend_error)?; Ok(objective_output( - format!("Edited objective {}", response.id), + format!("Edited objective {}", &response.resource_key), response, )?) } @@ -134,7 +139,7 @@ impl WorkspaceHttpObjectiveBackend { .await .map_err(backend_error)?; Ok(objective_output( - format!("Updated objective {} state", response.id), + format!("Updated objective {} state", &response.resource_key), response, )?) } @@ -142,6 +147,7 @@ impl WorkspaceHttpObjectiveBackend { async fn link_ticket(&self, input: ObjectiveLinkTicketInput) -> Result { let id = validate_id(&input.id, "ObjectiveLinkTicket")?; let ticket_id = validate_id(&input.ticket_id, "ObjectiveLinkTicket")?; + let ticket_resource_key = self.ticket_resource_key(ticket_id).await?; let url = format!("{}/ticket-links", self.objective_url(id)); let response = send_json::( self.client.as_ref(), @@ -154,7 +160,10 @@ impl WorkspaceHttpObjectiveBackend { .await .map_err(backend_error)?; Ok(objective_output( - format!("Linked ticket {ticket_id} to objective {}", response.id), + format!( + "Linked ticket {ticket_resource_key} to objective {}", + &response.resource_key + ), response, )?) } @@ -165,16 +174,46 @@ impl WorkspaceHttpObjectiveBackend { ) -> Result { let id = validate_id(&input.id, "ObjectiveUnlinkTicket")?; let ticket_id = validate_id(&input.ticket_id, "ObjectiveUnlinkTicket")?; + let ticket_resource_key = self.ticket_resource_key(ticket_id).await?; let url = format!("{}/ticket-links/{}", self.objective_url(id), ticket_id); let response = delete_json::(self.client.as_ref(), &url) .await .map_err(backend_error)?; Ok(objective_output( - format!("Unlinked ticket {ticket_id} from objective {}", response.id), + format!( + "Unlinked ticket {ticket_resource_key} from objective {}", + &response.resource_key + ), response, )?) } + async fn ticket_resource_key(&self, ticket_reference: &str) -> Result { + let workspace_id = self.client.workspace_id().unwrap_or_default(); + let response: serde_json::Value = decode_response( + self.client + .execute(WorkspaceRequest::get(format!( + "/api/w/{workspace_id}/tickets/{ticket_reference}" + ))) + .map_err(WorkspaceObjectiveBackendError::from) + .map_err(backend_error)?, + ) + .map_err(backend_error)?; + response + .get("resource_key") + .or_else(|| { + response + .get("meta") + .and_then(|meta| meta.get("resource_key")) + }) + .and_then(serde_json::Value::as_str) + .filter(|key| is_canonical_resource_key(key, "T-")) + .map(ToOwned::to_owned) + .ok_or_else(|| { + ToolError::ExecutionFailed("required T- human key is unavailable".to_string()) + }) + } + fn objective_url(&self, id: &str) -> String { let workspace_id = self.client.workspace_id().unwrap_or_default(); format!("/api/w/{workspace_id}/objectives/{id}") @@ -185,7 +224,7 @@ impl WorkspaceHttpObjectiveBackend { pub enum WorkspaceObjectiveBackendError { #[error("workspace objective backend request failed: {0}")] Request(#[from] crate::worker::WorkspaceClientError), - #[error("workspace objective backend returned HTTP {status}: {body}")] + #[error("workspace objective backend returned HTTP {status}")] Http { status: reqwest::StatusCode, body: String, @@ -247,10 +286,26 @@ fn decode_response Deserialize<'de>>( serde_json::from_str(&response.body).map_err(Into::into) } +fn is_canonical_resource_key(resource_key: &str, prefix: &str) -> bool { + resource_key.strip_prefix(prefix).is_some_and(|sequence| { + !sequence.is_empty() && sequence.bytes().all(|byte| byte.is_ascii_digit()) + }) +} + fn objective_output(summary: String, response: ObjectiveDetail) -> Result { + if !is_canonical_resource_key(&response.resource_key, "O-") { + return Err(ToolError::ExecutionFailed( + "required O- human key is unavailable".to_string(), + )); + } + let projected = serde_json::json!({ + "objective": &response.resource_key, + "title": response.title, + "state": response.state, + }); Ok(ToolOutput { summary, - content: Some(serde_json::to_string_pretty(&response).map_err(decode_error)?), + content: Some(serde_json::to_string_pretty(&projected).map_err(decode_error)?), attachments: Vec::new(), }) @@ -260,7 +315,7 @@ fn validate_id<'a>(id: &'a str, tool_name: &str) -> Result<&'a str, ToolError> { let id = id.trim(); if id.is_empty() || id.contains('/') { return Err(ToolError::InvalidArgument(format!( - "{tool_name} requires non-empty canonical id without '/'" + "{tool_name} requires a non-empty Objective reference without '/'" ))); } Ok(id) @@ -411,9 +466,9 @@ const EDIT_DESCRIPTION: &str = const SET_STATE_DESCRIPTION: &str = "Set an Objective state through Backend Workspace API authority."; const LINK_TICKET_DESCRIPTION: &str = - "Link a Ticket id to an Objective through Backend Workspace API authority."; + "Link a Ticket reference to an Objective through Backend Workspace API authority."; const UNLINK_TICKET_DESCRIPTION: &str = - "Unlink a Ticket id from an Objective through Backend Workspace API authority."; + "Unlink a Ticket reference from an Objective through Backend Workspace API authority."; fn list_schema() -> serde_json::Value { json!({ @@ -422,7 +477,7 @@ fn list_schema() -> serde_json::Value { "properties":{ "query":{"type":["string","null"]}, "states":{"type":"array","items":{"type":"string"},"default":[]}, - "linked_ticket_id":{"type":["string","null"]}, + "linked_ticket_id":{"type":["string","null"],"description":"Linked Ticket reference. Prefer T-*; canonical internal ids remain accepted for compatibility."}, "updated_after":{"type":["string","null"]}, "updated_before":{"type":["string","null"]}, "sort":{"type":["string","null"],"enum":["relevance","updated_desc","created_desc","title",null]}, @@ -438,7 +493,7 @@ fn show_schema() -> serde_json::Value { "additionalProperties": false, "required":["id"], "properties":{ - "id":{"type":"string"}, + "id":{"type":"string","description":"Objective reference. Prefer O-*; canonical internal ids remain accepted for compatibility."}, "event_limit":{"type":["integer","null"],"minimum":1,"maximum":50}, "event_cursor":{"type":["string","null"]} } @@ -454,7 +509,7 @@ fn create_schema() -> serde_json::Value { "title":{"type":"string","minLength":1}, "body_md":{"type":"string"}, "state":{"type":"string","default":"active"}, - "linked_tickets":{"type":"array","items":{"type":"string"}} + "linked_tickets":{"type":"array","items":{"type":"string"},"description":"Linked Ticket references. Prefer T-*; canonical internal ids remain accepted for compatibility."} } }) } @@ -465,7 +520,7 @@ fn edit_schema() -> serde_json::Value { "additionalProperties": false, "required":["id"], "properties":{ - "id":{"type":"string"}, + "id":{"type":"string","description":"Objective reference. Prefer O-*; canonical internal ids remain accepted for compatibility."}, "title":{"type":["string","null"]}, "old_string":{"type":["string","null"]}, "new_string":{"type":["string","null"]}, @@ -480,7 +535,7 @@ fn set_state_schema() -> serde_json::Value { "additionalProperties": false, "required":["id","state"], "properties":{ - "id":{"type":"string"}, + "id":{"type":"string","description":"Objective reference. Prefer O-*; canonical internal ids remain accepted for compatibility."}, "state":{"type":"string","minLength":1} } }) @@ -500,8 +555,8 @@ fn id_ticket_schema(required: &[&str]) -> serde_json::Value { "additionalProperties": false, "required": required, "properties":{ - "id":{"type":"string"}, - "ticket_id":{"type":"string"} + "id":{"type":"string","description":"Objective reference. Prefer O-*; canonical internal ids remain accepted for compatibility."}, + "ticket_id":{"type":"string","description":"Ticket reference. Prefer T-*; canonical internal ids remain accepted for compatibility."} } }) } @@ -595,21 +650,20 @@ fn default_state() -> String { #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)] struct ObjectiveDetail { - id: String, + resource_key: String, title: String, state: String, - created_at: Option, - updated_at: Option, - linked_tickets: Vec, - body: String, - body_truncated: bool, - record_source: String, } #[cfg(test)] mod tests { use super::*; use agen::tool::ToolDefinition; + use std::{ + io::{Read, Write}, + net::TcpListener, + thread, + }; fn tool_names(definitions: Vec) -> Vec { let mut names = definitions @@ -656,4 +710,135 @@ mod tests { let link = link_ticket_schema(); assert_eq!(link["required"], json!(["id", "ticket_id"])); } + + #[tokio::test(flavor = "multi_thread")] + async fn objective_show_summary_uses_projected_human_key() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let server = thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut buffer = [0_u8; 8192]; + let len = stream.read(&mut buffer).unwrap(); + let request = String::from_utf8_lossy(&buffer[..len]); + assert!( + request.starts_with("POST /api/w/workspace/objectives/00001INTERNAL/show HTTP/1.1") + ); + let body = serde_json::json!({ + "id": "00001INTERNAL", + "resource_key": "O-3", + "title": "Objective", + "body": "Body", + "state": "active", + "created_at": null, + "updated_at": null, + "linked_ticket_summaries": [], + "events": [], + "event_page": {"next_cursor": null, "has_more": false} + }) + .to_string(); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", + body.len(), + body + ) + .unwrap(); + }); + let backend = WorkspaceHttpObjectiveBackend::new(Arc::new( + crate::worker::TestWorkspaceHttpClient::new("workspace", base_url), + )); + + let output = backend + .show(ShowObjectiveInput { + id: "00001INTERNAL".to_string(), + event_limit: None, + event_cursor: None, + }) + .await + .unwrap(); + + server.join().unwrap(); + assert_eq!(output.summary, "Read objective O-3"); + assert!(!output.summary.contains("00001INTERNAL")); + assert!(!output.content.unwrap().contains("00001INTERNAL")); + } + + #[tokio::test(flavor = "multi_thread")] + async fn objective_link_summaries_resolve_internal_ticket_ids_to_human_keys() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let server = thread::spawn(move || { + for mutation in ["POST", "DELETE"] { + let (mut stream, _) = listener.accept().unwrap(); + let mut buffer = [0_u8; 8192]; + let len = stream.read(&mut buffer).unwrap(); + let request = String::from_utf8_lossy(&buffer[..len]); + assert!(request.starts_with("GET /api/w/workspace/tickets/00001INTERNAL HTTP/1.1")); + let response_body = serde_json::json!({"resource_key": "T-7"}).to_string(); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", + response_body.len(), + response_body + ) + .unwrap(); + + let (mut stream, _) = listener.accept().unwrap(); + let mut buffer = [0_u8; 8192]; + let len = stream.read(&mut buffer).unwrap(); + let request = String::from_utf8_lossy(&buffer[..len]); + assert!(request.starts_with(&format!( + "{mutation} /api/w/workspace/objectives/O-3/ticket-links" + ))); + let response_body = serde_json::json!({ + "resource_key": "O-3", + "title": "Objective", + "state": "active" + }) + .to_string(); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", + response_body.len(), + response_body + ) + .unwrap(); + } + }); + let backend = WorkspaceHttpObjectiveBackend::new(Arc::new( + crate::worker::TestWorkspaceHttpClient::new("workspace", base_url), + )); + + let linked = backend + .link_ticket(ObjectiveLinkTicketInput { + id: "O-3".to_string(), + ticket_id: "00001INTERNAL".to_string(), + }) + .await + .unwrap(); + let unlinked = backend + .unlink_ticket(ObjectiveUnlinkTicketInput { + id: "O-3".to_string(), + ticket_id: "00001INTERNAL".to_string(), + }) + .await + .unwrap(); + + server.join().unwrap(); + for output in [linked, unlinked] { + assert!(output.summary.contains("T-7")); + assert!(!output.summary.contains("00001INTERNAL")); + assert!(!output.content.unwrap().contains("00001INTERNAL")); + } + } + + #[test] + fn objective_output_rejects_noncanonical_human_keys() { + let response = ObjectiveDetail { + resource_key: "O-internal".to_string(), + title: "Objective".to_string(), + state: "active".to_string(), + }; + assert!(objective_output("created".to_string(), response).is_err()); + } } diff --git a/crates/worker/src/feature/builtin/resource_projection.rs b/crates/worker/src/feature/builtin/resource_projection.rs new file mode 100644 index 00000000..0eb32758 --- /dev/null +++ b/crates/worker/src/feature/builtin/resource_projection.rs @@ -0,0 +1,794 @@ +use serde::Serialize; +use serde_json::{Map, Value}; + +#[derive(Debug, Serialize)] +pub(super) struct ModelTicketQueryResponse { + tickets: Vec, + next_cursor: Option, + has_more: bool, +} + +#[derive(Debug, Serialize)] +struct ModelTicketQueryItem { + ticket: String, + title: String, + state: String, + readiness: Option, + priority: Option, + created_at: Option, + updated_at: Option, + workspace_action_priority: Option, + matched_fields: Vec, + snippet: Option, + current_coder: Option, + linked_objectives: Vec, + relation_count: usize, + blocker_count: usize, + unresolved_blocker_count: usize, + unresolved_review_count: usize, + evidence: Option, + merge_request: Option, +} + +#[derive(Debug, Serialize)] +pub(super) struct ModelTicketDetail { + ticket: String, + title: String, + body: String, + state: String, + readiness: Option, + priority: Option, + created_at: Option, + updated_at: Option, + thread: Vec, + relations: ModelTicketRelations, + linked_objectives: Vec, + assignments: Vec, + current_coder: Option, + implementation_reports: Vec, + merge_request: Option, + evidence: Option, + actions: Option, + event_page: Option, +} + +#[derive(Debug, Serialize)] +pub(super) struct ModelObjectiveQueryResponse { + objectives: Vec, + next_cursor: Option, + has_more: bool, +} + +#[derive(Debug, Serialize)] +struct ModelObjectiveQueryItem { + objective: String, + title: String, + summary: Option, + state: String, + created_at: Option, + updated_at: Option, + linked_tickets: Vec, + linked_ticket_count: usize, +} + +#[derive(Debug, Serialize)] +pub(super) struct ModelObjectiveDetail { + objective: String, + title: String, + body: String, + state: String, + created_at: Option, + updated_at: Option, + linked_tickets: Vec, + events: Vec, + event_page: ModelObjectiveEventPage, +} + +impl ModelObjectiveDetail { + pub(super) fn objective_ref(&self) -> &str { + &self.objective + } +} + +#[derive(Debug, Serialize)] +struct ModelWorkerSummary { + worker: String, +} + +#[derive(Debug, Serialize)] +struct ModelTicketEvent { + sequence: usize, + kind: String, + body: Option, + created_at: Option, +} + +#[derive(Debug, Serialize, Default)] +struct ModelTicketRelations { + outgoing: Vec, + incoming: Vec, + blockers: Vec, + notices: Vec, +} + +#[derive(Debug, Serialize)] +struct ModelRelation { + ticket: String, + kind: String, + note: Option, + created_at: Option, +} + +#[derive(Debug, Serialize)] +struct ModelBlocker { + ticket: String, + kind: String, + state: Option, + resolved: bool, +} + +#[derive(Debug, Serialize)] +struct ModelNotice { + kind: String, +} + +#[derive(Debug, Serialize)] +struct ModelObjectiveSummary { + objective: String, + title: String, + state: String, +} + +#[derive(Debug, Serialize)] +struct ModelTicketSummary { + ticket: String, + title: String, + state: String, +} + +#[derive(Debug, Serialize)] +struct ModelAssignment { + role: String, + principal: String, + assigned_at: String, +} + +#[derive(Debug, Serialize)] +struct ModelEvidenceEvent { + sequence: usize, + kind: String, + created_at: Option, + excerpt: String, +} + +#[derive(Debug, Serialize)] +struct ModelMergeRequest { + state: String, + selector_from: Option, + selector_to: String, + review_status: String, + subject_ref: Option, + review_excerpt: Option, +} + +#[derive(Debug, Serialize)] +struct ModelTicketEvidence { + has_merge_request: bool, + has_current_subject_ref: bool, + has_review_request: bool, + has_commit: bool, + review_status: Option, + approved_current_subject: bool, + unresolved_request_changes: bool, + complete_for_integration: bool, + missing: Vec, +} + +#[derive(Debug, Serialize)] +struct ModelTicketActions { + can_assign_orchestrator: bool, + can_unassign_orchestrator: bool, + can_queue: bool, + can_start_manual_coder: bool, +} + +#[derive(Debug, Serialize)] +struct ModelEventPage { + next_cursor: Option, + has_more: bool, +} + +#[derive(Debug, Serialize)] +struct ModelObjectiveEvent { + kind: String, + created_at: String, + body: Option, +} + +#[derive(Debug, Serialize)] +struct ModelObjectiveEventPage { + next_cursor: Option, + has_more: bool, +} + +pub(super) fn project_ticket_query(value: Value) -> Result { + let root = object(&value, "Ticket query response")?; + let page = object_field(root, "page")?; + let tickets = array_field(root, "items")? + .iter() + .map(project_ticket_query_item) + .collect::, _>>()?; + Ok(ModelTicketQueryResponse { + tickets, + next_cursor: optional_string(page, "next_cursor")?, + has_more: bool_field(page, "has_more")?, + }) +} + +fn project_ticket_query_item(value: &Value) -> Result { + let item = object(value, "Ticket query item")?; + Ok(ModelTicketQueryItem { + ticket: human_ref(item, "resource_key", "T-")?, + title: string_field(item, "title")?, + state: string_field(item, "state")?, + readiness: optional_string(item, "readiness")?, + priority: optional_string(item, "priority")?, + created_at: optional_string(item, "created_at")?, + updated_at: optional_string(item, "updated_at")?, + workspace_action_priority: optional_string(item, "workspace_action_priority")?, + matched_fields: string_array(item, "matched_fields")?, + snippet: optional_string(item, "snippet")?, + current_coder: item + .get("current_coder") + .filter(|value| !value.is_null()) + .map(project_worker) + .transpose()?, + linked_objectives: string_array(item, "linked_objective_keys")? + .into_iter() + .map(|key| validate_human_ref(key, "O-")) + .collect::, _>>()?, + relation_count: usize_field(item, "relation_count")?, + blocker_count: usize_field(item, "blocker_count")?, + unresolved_blocker_count: usize_field(item, "unresolved_blocker_count")?, + unresolved_review_count: usize_field(item, "unresolved_review_count")?, + evidence: item.get("evidence").map(project_evidence).transpose()?, + merge_request: item + .get("merge_request") + .filter(|value| !value.is_null()) + .map(project_merge_request) + .transpose()?, + }) +} + +pub(super) fn project_ticket_detail(value: Value) -> Result { + let root = object(&value, "Ticket detail response")?; + let current_coder = root + .get("current_coder") + .filter(|value| !value.is_null()) + .map(project_worker) + .transpose()?; + let assignments = array_field(root, "assignments")? + .iter() + .map(|assignment| project_assignment(assignment, current_coder.as_ref())) + .collect::, _>>()?; + + Ok(ModelTicketDetail { + ticket: human_ref(root, "resource_key", "T-")?, + title: string_field(root, "title")?, + body: string_field(root, "body")?, + state: string_field(root, "state")?, + readiness: optional_string(root, "readiness")?, + priority: optional_string(root, "priority")?, + created_at: optional_string(root, "created_at")?, + updated_at: optional_string(root, "updated_at")?, + thread: array_field(root, "events")? + .iter() + .map(project_ticket_event) + .collect::, _>>()?, + relations: project_relations(root.get("relations"))?, + linked_objectives: array_field(root, "linked_objectives")? + .iter() + .map(project_objective_summary) + .collect::, _>>()?, + assignments, + current_coder, + implementation_reports: array_field(root, "implementation_reports")? + .iter() + .map(project_evidence_event) + .collect::, _>>()?, + merge_request: root + .get("merge_request") + .filter(|value| !value.is_null()) + .map(project_merge_request) + .transpose()?, + evidence: root.get("evidence").map(project_evidence).transpose()?, + actions: root + .get("action_eligibility") + .filter(|value| !value.is_null()) + .map(project_actions) + .transpose()?, + event_page: root + .get("event_page") + .filter(|value| !value.is_null()) + .map(project_event_page) + .transpose()?, + }) +} + +pub(super) fn project_objective_query(value: Value) -> Result { + let root = object(&value, "Objective query response")?; + let page = object_field(root, "page")?; + Ok(ModelObjectiveQueryResponse { + objectives: array_field(root, "items")? + .iter() + .map(project_objective_query_item) + .collect::, _>>()?, + next_cursor: optional_string(page, "next_cursor")?, + has_more: bool_field(page, "has_more")?, + }) +} + +fn project_objective_query_item(value: &Value) -> Result { + let item = object(value, "Objective query item")?; + let linked_tickets = string_array(item, "linked_ticket_keys")? + .into_iter() + .map(|key| validate_human_ref(key, "T-")) + .collect::, _>>()?; + Ok(ModelObjectiveQueryItem { + objective: human_ref(item, "resource_key", "O-")?, + title: string_field(item, "title")?, + summary: optional_string(item, "snippet")?, + state: string_field(item, "state")?, + created_at: optional_string(item, "created_at")?, + updated_at: optional_string(item, "updated_at")?, + linked_ticket_count: linked_tickets.len(), + linked_tickets, + }) +} + +pub(super) fn project_objective_detail(value: Value) -> Result { + let root = object(&value, "Objective detail response")?; + Ok(ModelObjectiveDetail { + objective: human_ref(root, "resource_key", "O-")?, + title: string_field(root, "title")?, + body: string_field(root, "body")?, + state: string_field(root, "state")?, + created_at: optional_string(root, "created_at")?, + updated_at: optional_string(root, "updated_at")?, + linked_tickets: array_field(root, "linked_ticket_summaries")? + .iter() + .map(project_ticket_summary) + .collect::, _>>()?, + events: array_field(root, "events")? + .iter() + .map(project_objective_event) + .collect::, _>>()?, + event_page: project_objective_event_page( + root.get("event_page") + .ok_or_else(|| "Objective detail response is missing event_page".to_string())?, + )?, + }) +} + +fn project_worker(value: &Value) -> Result { + let worker = object(value, "Worker summary")?; + Ok(ModelWorkerSummary { + worker: human_ref(worker, "worker_resource_key", "W-")?, + }) +} + +fn project_ticket_event(value: &Value) -> Result { + let event = object(value, "Ticket event")?; + Ok(ModelTicketEvent { + sequence: usize_field(event, "sequence")?, + kind: string_field(event, "kind")?, + body: match event.get("body") { + None | Some(Value::Null) => None, + Some(Value::String(body)) => Some(body.clone()), + Some(_) => return Err("invalid Ticket event body".to_string()), + }, + created_at: optional_string(event, "at")?, + }) +} + +fn project_relations(value: Option<&Value>) -> Result { + let Some(value) = value else { + return Ok(ModelTicketRelations::default()); + }; + let relations = object(value, "Ticket relations")?; + Ok(ModelTicketRelations { + outgoing: array_field(relations, "outgoing")? + .iter() + .map(|value| project_relation(value, "target_resource_key", "kind")) + .collect::, _>>()?, + incoming: array_field(relations, "incoming")? + .iter() + .map(|value| project_relation(value, "source_resource_key", "forward_kind")) + .collect::, _>>()?, + blockers: array_field(relations, "blockers")? + .iter() + .map(project_blocker) + .collect::, _>>()?, + notices: array_field(relations, "notices")? + .iter() + .map(project_notice) + .collect::, _>>()?, + }) +} + +fn project_relation( + value: &Value, + ticket_key: &str, + kind_key: &str, +) -> Result { + let relation = object(value, "Ticket relation")?; + let relation_data = relation.get("relation").and_then(Value::as_object); + let kind = if kind_key == "kind" { + relation_data + .ok_or_else(|| "Ticket relation is missing relation data".to_string()) + .and_then(|data| string_field(data, "kind"))? + } else { + string_field(relation, kind_key)? + }; + let note = match relation_data { + Some(data) => optional_string(data, "note")?, + None => optional_string(relation, "note")?, + }; + let created_at = match relation_data { + Some(data) => optional_string(data, "at")?, + None => optional_string(relation, "at")?, + }; + Ok(ModelRelation { + ticket: human_ref(relation, ticket_key, "T-")?, + kind, + note, + created_at, + }) +} + +fn project_blocker(value: &Value) -> Result { + let blocker = object(value, "Ticket blocker")?; + Ok(ModelBlocker { + ticket: human_ref(blocker, "blocking_resource_key", "T-")?, + kind: string_field(blocker, "relation_kind")?, + state: optional_string(blocker, "blocking_state")?, + resolved: bool_field(blocker, "resolved")?, + }) +} + +fn project_notice(value: &Value) -> Result { + let notice = object(value, "Ticket notice")?; + Ok(ModelNotice { + kind: string_field(notice, "kind")?, + }) +} + +fn project_objective_summary(value: &Value) -> Result { + let summary = object(value, "Objective summary")?; + Ok(ModelObjectiveSummary { + objective: human_ref(summary, "resource_key", "O-")?, + title: string_field(summary, "title")?, + state: string_field(summary, "state")?, + }) +} + +fn project_ticket_summary(value: &Value) -> Result { + let summary = object(value, "Ticket summary")?; + Ok(ModelTicketSummary { + ticket: human_ref(summary, "resource_key", "T-")?, + title: string_field(summary, "title")?, + state: string_field(summary, "state")?, + }) +} + +fn project_assignment( + value: &Value, + current_coder: Option<&ModelWorkerSummary>, +) -> Result { + let assignment = object(value, "Ticket assignment")?; + let principal = object_field(assignment, "principal")?; + let kind = string_field(principal, "kind")?; + let principal = match kind.as_str() { + "worker" => current_coder + .map(|coder| coder.worker.clone()) + .ok_or_else(|| { + "Worker assignment is missing a Workspace human key projection".to_string() + })?, + "workspace_agent" => format!("workspace-agent:{}", string_field(principal, "agent_key")?), + "user" => "user".to_string(), + other => format!("source:{other}"), + }; + Ok(ModelAssignment { + role: string_field(assignment, "role")?, + principal, + assigned_at: string_field(assignment, "assigned_at")?, + }) +} + +fn project_evidence_event(value: &Value) -> Result { + let event = object(value, "Ticket evidence event")?; + Ok(ModelEvidenceEvent { + sequence: usize_field(event, "sequence")?, + kind: string_field(event, "kind")?, + created_at: optional_string(event, "at")?, + excerpt: string_field(event, "excerpt")?, + }) +} + +fn project_merge_request(value: &Value) -> Result { + let merge = object(value, "Merge Request summary")?; + Ok(ModelMergeRequest { + state: string_field(merge, "state")?, + selector_from: optional_string(merge, "selector_from")?, + selector_to: string_field(merge, "selector_to")?, + review_status: string_field(merge, "review_status")?, + subject_ref: optional_string(merge, "subject_ref")?, + review_excerpt: optional_string(merge, "review_excerpt")?, + }) +} + +fn project_evidence(value: &Value) -> Result { + let evidence = object(value, "Ticket evidence")?; + Ok(ModelTicketEvidence { + has_merge_request: bool_field(evidence, "has_merge_request")?, + has_current_subject_ref: bool_field(evidence, "has_current_subject_ref")?, + has_review_request: bool_field(evidence, "has_review_request")?, + has_commit: bool_field(evidence, "has_commit")?, + review_status: optional_string(evidence, "review_status")?, + approved_current_subject: bool_field(evidence, "approved_current_subject")?, + unresolved_request_changes: bool_field(evidence, "unresolved_request_changes")?, + complete_for_integration: bool_field(evidence, "complete_for_integration")?, + missing: string_array(evidence, "missing")?, + }) +} + +fn project_actions(value: &Value) -> Result { + let actions = object(value, "Ticket actions")?; + Ok(ModelTicketActions { + can_assign_orchestrator: bool_field(actions, "can_assign_orchestrator")?, + can_unassign_orchestrator: bool_field(actions, "can_unassign_orchestrator")?, + can_queue: bool_field(actions, "can_queue")?, + can_start_manual_coder: bool_field(actions, "can_start_manual_coder")?, + }) +} + +fn project_event_page(value: &Value) -> Result { + let page = object(value, "Ticket event page")?; + Ok(ModelEventPage { + next_cursor: optional_string(page, "next_cursor")?, + has_more: bool_field(page, "has_more")?, + }) +} + +fn project_objective_event(value: &Value) -> Result { + let event = object(value, "Objective event")?; + let body = optional_string(event, "body")?; + Ok(ModelObjectiveEvent { + kind: string_field(event, "kind")?, + created_at: string_field(event, "created_at")?, + body, + }) +} + +fn project_objective_event_page(value: &Value) -> Result { + let page = object(value, "Objective event page")?; + Ok(ModelObjectiveEventPage { + next_cursor: optional_string(page, "next_cursor")?, + has_more: bool_field(page, "has_more")?, + }) +} + +fn object<'a>(value: &'a Value, context: &str) -> Result<&'a Map, String> { + value + .as_object() + .ok_or_else(|| format!("{context} must be an object")) +} + +fn object_field<'a>( + object: &'a Map, + key: &str, +) -> Result<&'a Map, String> { + object + .get(key) + .and_then(Value::as_object) + .ok_or_else(|| format!("missing or invalid {key}")) +} + +fn array_field<'a>(object: &'a Map, key: &str) -> Result<&'a [Value], String> { + object + .get(key) + .and_then(Value::as_array) + .map(Vec::as_slice) + .ok_or_else(|| format!("missing or invalid {key}")) +} + +fn string_field(object: &Map, key: &str) -> Result { + object + .get(key) + .and_then(Value::as_str) + .map(ToOwned::to_owned) + .ok_or_else(|| format!("missing or invalid {key}")) +} + +fn optional_string(object: &Map, key: &str) -> Result, String> { + match object.get(key) { + None | Some(Value::Null) => Ok(None), + Some(Value::String(value)) => Ok(Some(value.clone())), + Some(_) => Err(format!("invalid {key}")), + } +} + +fn bool_field(object: &Map, key: &str) -> Result { + object + .get(key) + .and_then(Value::as_bool) + .ok_or_else(|| format!("missing or invalid {key}")) +} + +fn usize_field(object: &Map, key: &str) -> Result { + object + .get(key) + .and_then(Value::as_u64) + .and_then(|value| usize::try_from(value).ok()) + .ok_or_else(|| format!("missing or invalid {key}")) +} + +fn string_array(object: &Map, key: &str) -> Result, String> { + array_field(object, key)? + .iter() + .map(|value| { + value + .as_str() + .map(ToOwned::to_owned) + .ok_or_else(|| format!("invalid {key}")) + }) + .collect() +} + +fn human_ref(object: &Map, key: &str, prefix: &str) -> Result { + let value = object + .get(key) + .and_then(Value::as_str) + .map(ToOwned::to_owned) + .ok_or_else(|| format!("required {prefix} human key is unavailable"))?; + validate_human_ref(value, prefix) +} + +fn validate_human_ref(value: String, prefix: &str) -> Result { + let valid = value.strip_prefix(prefix).is_some_and(|sequence| { + !sequence.is_empty() && sequence.bytes().all(|byte| byte.is_ascii_digit()) + }); + if valid { + Ok(value) + } else { + Err(format!("required {prefix} human key is unavailable")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn objective_projection_exposes_only_human_resource_references() { + let projected = project_objective_detail(json!({ + "id": "00001M10HW6BV", + "resource_key": "O-543", + "title": "Objective", + "body": "Body", + "state": "active", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-02T00:00:00Z", + "linked_tickets": ["00001M0E82D1V"], + "linked_ticket_summaries": [{ + "id": "00001M0E82D1V", + "resource_key": "T-496", + "title": "Ticket", + "state": "done", + "updated_at": "2026-01-02T00:00:00Z" + }], + "events": [{ + "sequence": 3, + "event_ref": "objective-event-3", + "kind": "linked_ticket", + "created_at": "2026-01-02T00:00:00Z", + "body": "linked" + }], + "event_page": {"next_cursor": null, "has_more": false, "window_start_sequence": 3, "window_end_sequence": 3} + })).expect("projection"); + let json = serde_json::to_value(projected).expect("serialize"); + let text = json.to_string(); + assert!(text.contains("O-543")); + assert!(text.contains("T-496")); + assert!(!text.contains("00001M10HW6BV")); + assert!(!text.contains("00001M0E82D1V")); + assert!(!text.contains("event_ref")); + } + + #[test] + fn query_projections_accept_workspace_api_shapes_and_scrub_internal_ids() { + let ticket = project_ticket_query(json!({ + "page": {"next_cursor": null, "has_more": false}, + "record_authority": "workspace_sqlite", + "items": [{ + "id": "00001TICKETINTERNAL", + "resource_key": "T-543", + "title": "Ticket", + "state": "inprogress", + "readiness": null, + "priority": "high", + "created_at": null, + "updated_at": "2026-01-01T00:00:00Z", + "workspace_action_priority": "active_work", + "matched_fields": ["title"], + "snippet": "Ticket", + "current_coder": {"runtime_id": "runtime-internal", "worker_id": "worker-internal", "worker_resource_key": "W-12"}, + "linked_objective_ids": ["00001OBJECTIVEINTERNAL"], + "linked_objective_keys": ["O-6"], + "relation_count": 0, + "blocker_count": 0, + "unresolved_blocker_count": 0, + "unresolved_review_count": 0, + "evidence": { + "has_merge_request": false, + "has_current_subject_ref": false, + "has_review_request": false, + "has_commit": false, + "review_status": null, + "approved_current_subject": false, + "unresolved_request_changes": false, + "complete_for_integration": false, + "missing": ["merge_request"] + }, + "merge_request": null + }] + })).expect("Ticket query projection"); + let ticket_json = serde_json::to_string(&ticket).expect("serialize Ticket query"); + assert!(ticket_json.contains("T-543")); + assert!(ticket_json.contains("O-6")); + assert!(ticket_json.contains("W-12")); + assert!(!ticket_json.contains("00001TICKETINTERNAL")); + assert!(!ticket_json.contains("runtime-internal")); + assert!(!ticket_json.contains("worker-internal")); + + let objective = project_objective_query(json!({ + "page": {"next_cursor": null, "has_more": false}, + "record_authority": "workspace_sqlite", + "items": [{ + "id": "00001OBJECTIVEINTERNAL", + "resource_key": "O-6", + "title": "Objective", + "state": "active", + "created_at": null, + "updated_at": null, + "matched_fields": [], + "snippet": null, + "linked_ticket_count": 1, + "linked_tickets": ["00001TICKETINTERNAL"], + "linked_ticket_keys": ["T-543"] + }] + })) + .expect("Objective query projection"); + let objective_json = serde_json::to_string(&objective).expect("serialize Objective query"); + assert!(objective_json.contains("O-6")); + assert!(objective_json.contains("T-543")); + assert!(objective_json.contains("\"summary\":null")); + assert!(!objective_json.contains("00001OBJECTIVEINTERNAL")); + assert!(!objective_json.contains("00001TICKETINTERNAL")); + } + + #[test] + fn human_resource_projection_rejects_noncanonical_keys() { + for (key, prefix) in [("T-key", "T-"), ("O-", "O-"), ("W-1x", "W-")] { + assert!(validate_human_ref(key.to_string(), prefix).is_err()); + } + } + + #[test] + fn ticket_projection_fails_closed_without_worker_resource_key() { + let error = project_worker(&json!({"worker_resource_key": null})) + .expect_err("missing W-key must fail"); + assert!(error.contains("W-")); + } +} diff --git a/crates/worker/src/feature/builtin/session_explore.rs b/crates/worker/src/feature/builtin/session_explore.rs index 24e541ed..09d5bc24 100644 --- a/crates/worker/src/feature/builtin/session_explore.rs +++ b/crates/worker/src/feature/builtin/session_explore.rs @@ -193,6 +193,7 @@ impl Tool for ShowOverviewTool { .map(|entry| { serde_json::json!({ "entry_ref": entry.id, + "origin": entry.origin, "entry_range": entry.entry_range, "kind": entry.kind.as_str(), "label": entry.label, @@ -234,15 +235,16 @@ impl Tool for SearchEntriesTool { .transpose()?; let from = params.from.as_deref().map(parse_entry_ref).transpose()?; let through = params.through.as_deref().map(parse_entry_ref).transpose()?; + let view = self.state.view(); if let (Some(from), Some(through)) = (&from, &through) { - if from.source_index() > through.source_index() { + if view.source_index_for_ref(from) > view.source_index_for_ref(through) { return Err(ToolError::InvalidArgument( "SearchEntries from must not be after through".to_string(), )); } } let limit = bounded_limit(params.limit, DEFAULT_PAGE_LIMIT, MAX_PAGE_LIMIT); - let hits = self.state.view().search(&SearchOptions { + let hits = view.search(&SearchOptions { query: params.query, kind, tool_part, @@ -318,6 +320,7 @@ impl Tool for ReadEntryTool { .map(|entry| { serde_json::json!({ "entry_ref": entry.id, + "origin": entry.origin, "entry_range": entry.entry_range, "kind": entry.kind.as_str(), "tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()), diff --git a/crates/worker/src/feature/builtin/ticket.rs b/crates/worker/src/feature/builtin/ticket.rs index ae39e3cb..0194d9d0 100644 --- a/crates/worker/src/feature/builtin/ticket.rs +++ b/crates/worker/src/feature/builtin/ticket.rs @@ -33,6 +33,8 @@ use crate::feature::{ use crate::worker::{WorkspaceClient, WorkspaceRequest, WorkspaceRequestMethod}; use agen::tool::{Tool, ToolError, ToolExecutionContext, ToolMeta, ToolOutput}; +use super::resource_projection::{project_ticket_detail, project_ticket_query}; + #[derive(Clone, Copy)] enum WorkspaceTicketReadKind { Query, @@ -153,8 +155,10 @@ struct WorkspaceQueryTicketInput { /// stale_after_rescope, and missing_evidence. #[serde(default)] attention: Vec, + /// Related Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. related_ticket_id: Option, relation_kind: Option, + /// Linked Objective reference. Prefer `O-*`; canonical internal ids remain accepted for compatibility. linked_objective_id: Option, updated_after: Option, updated_before: Option, @@ -169,6 +173,7 @@ struct WorkspaceQueryTicketInput { #[derive(Debug, Deserialize, Serialize, JsonSchema)] struct WorkspaceShowTicketInput { + /// Ticket reference. Prefer `T-*`; canonical internal ids remain accepted for compatibility. id: String, /// Most-recent thread entries to return, bounded by the Backend to 1..=50. event_limit: Option, @@ -229,13 +234,27 @@ impl Tool for WorkspaceTicketReadTool { .map_err(|error| ToolError::ExecutionFailed(error.to_string()))?; if !response.is_success() { return Err(ToolError::ExecutionFailed(format!( - "Workspace Ticket API returned HTTP {}: {}", - response.status, response.body + "Workspace Ticket API request failed with HTTP status {}", + response.status ))); } + let response_value: Value = serde_json::from_str(&response.body).map_err(|error| { + ToolError::ExecutionFailed(format!( + "Workspace Ticket API returned invalid JSON: {error}" + )) + })?; + let content = match self.kind { + WorkspaceTicketReadKind::Query => serde_json::to_string( + &project_ticket_query(response_value).map_err(ToolError::ExecutionFailed)?, + ), + WorkspaceTicketReadKind::Show => serde_json::to_string( + &project_ticket_detail(response_value).map_err(ToolError::ExecutionFailed)?, + ), + } + .map_err(|error| ToolError::Internal(error.to_string()))?; Ok(ToolOutput { summary: self.kind.name().to_string(), - content: Some(response.body), + content: Some(content), attachments: Vec::new(), }) } @@ -733,14 +752,69 @@ impl WorkspaceHttpTicketBackend { })?; if !response.is_success() { return Err(TicketError::Conflict(format!( - "ticket REST API returned HTTP {}: {}", - response.status, response.body + "ticket REST API request failed with HTTP status {}", + response.status ))); } - serde_json::from_str(&response.body) + let mut value: Value = serde_json::from_str(&response.body).map_err(|error| { + TicketError::Conflict(format!("decode ticket REST response: {error}")) + })?; + Self::canonicalize_ticket_references(&mut value); + serde_json::from_value(value) .map_err(|error| TicketError::Conflict(format!("decode ticket REST response: {error}"))) } + fn canonicalize_ticket_references(value: &mut Value) { + match value { + Value::Array(values) => { + for value in values { + Self::canonicalize_ticket_references(value); + } + } + Value::Object(object) => { + for value in object.values_mut() { + Self::canonicalize_ticket_references(value); + } + if let Some(resource_key) = object + .get("resource_key") + .and_then(Value::as_str) + .filter(|key| is_canonical_ticket_resource_key(key)) + .map(ToOwned::to_owned) + && object.contains_key("id") + { + object.insert("id".to_string(), Value::String(resource_key)); + } + } + _ => {} + } + } + + fn resolve_ticket_resource_key( + client: Arc, + base: &str, + reference: &TicketIdOrSlug, + ) -> TicketResult { + let response: Value = Self::request( + client, + WorkspaceRequestMethod::Get, + format!("{base}/{}", Self::ticket_path(reference)), + None, + )?; + response + .get("resource_key") + .or_else(|| { + response + .get("meta") + .and_then(|meta| meta.get("resource_key")) + }) + .and_then(Value::as_str) + .filter(|key| is_canonical_ticket_resource_key(key)) + .map(ToOwned::to_owned) + .ok_or_else(|| { + TicketError::Conflict("required Ticket human key is unavailable".to_string()) + }) + } + fn request_unit( client: Arc, method: WorkspaceRequestMethod, @@ -760,8 +834,8 @@ impl WorkspaceHttpTicketBackend { })?; if !response.is_success() { return Err(TicketError::Conflict(format!( - "ticket REST API returned HTTP {}: {}", - response.status, response.body + "ticket REST API request failed with HTTP status {}", + response.status ))); } Ok(TicketBackendOperationResult::Unit) @@ -802,12 +876,22 @@ impl WorkspaceHttpTicketBackend { Ok(TicketBackendOperationResult::Tickets(tickets)) } TicketBackendOperation::Show { id } => { - let ticket = Self::request( + let ticket: Ticket = Self::request( client, WorkspaceRequestMethod::Get, format!("{base}/{}/record", Self::ticket_path(&id)), None, )?; + if !ticket + .meta + .resource_key + .as_deref() + .is_some_and(is_canonical_ticket_resource_key) + { + return Err(TicketError::Conflict( + "required Ticket human key is unavailable".to_string(), + )); + } Ok(TicketBackendOperationResult::Ticket(ticket)) } TicketBackendOperation::Create { input } => { @@ -910,7 +994,14 @@ impl WorkspaceHttpTicketBackend { })?), ), TicketBackendOperation::AddTicketRelation { id, relation } => { - let relation = Self::request( + let source_resource_key = + Self::resolve_ticket_resource_key(client.clone(), &base, &id)?; + let target_resource_key = Self::resolve_ticket_resource_key( + client.clone(), + &base, + &TicketIdOrSlug::Id(relation.target.clone()), + )?; + let mut relation: TicketRelation = Self::request( client, WorkspaceRequestMethod::Post, format!("{base}/{}/relations", Self::ticket_path(&id)), @@ -918,20 +1009,30 @@ impl WorkspaceHttpTicketBackend { TicketError::Conflict(format!("serialize Ticket relation: {error}")) })?), )?; + relation.ticket_id = source_resource_key; + relation.target = target_resource_key; + relation.author = "workspace".to_string(); Ok(TicketBackendOperationResult::Relation(relation)) } TicketBackendOperation::RemoveTicketRelation { id, kind, target } => { + let source_resource_key = + Self::resolve_ticket_resource_key(client.clone(), &base, &id)?; + let target_resource_key = + Self::resolve_ticket_resource_key(client.clone(), &base, &target)?; let target = match target { TicketIdOrSlug::Id(value) | TicketIdOrSlug::Slug(value) | TicketIdOrSlug::Query(value) => value, }; - let relation = Self::request( + let mut relation: TicketRelation = Self::request( client, WorkspaceRequestMethod::Delete, format!("{base}/{}/relations", Self::ticket_path(&id)), Some(serde_json::json!({ "kind": kind, "target": target })), )?; + relation.ticket_id = source_resource_key; + relation.target = target_resource_key; + relation.author = "workspace".to_string(); Ok(TicketBackendOperationResult::Relation(relation)) } TicketBackendOperation::QueryTicketRelations { ticket, kind } => { @@ -1266,6 +1367,23 @@ mod tests { .expect("tool exists") } + #[test] + fn workspace_ticket_backend_canonicalizes_model_facing_ticket_ids() { + let mut value = serde_json::json!({ + "id": "00001INTERNAL", + "resource_key": "T-42", + "nested": { + "id": "00002INTERNAL", + "resource_key": "T-43" + }, + "body": "user-authored 00003BODY stays unchanged" + }); + WorkspaceHttpTicketBackend::canonicalize_ticket_references(&mut value); + assert_eq!(value["id"], "T-42"); + assert_eq!(value["nested"]["id"], "T-43"); + assert_eq!(value["body"], "user-authored 00003BODY stays unchanged"); + } + #[test] fn workspace_ticket_reads_expose_bounded_query_and_show_contracts_without_legacy_aliases() { let client: Arc = Arc::new( @@ -1742,11 +1860,102 @@ provider = "github" server.join().unwrap(); } + #[test] + fn workspace_http_backend_records_relation_with_authoritative_human_keys() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let server = thread::spawn(move || { + for (expected_path, resource_key) in [ + ("GET /api/w/workspace-a/tickets/01SOURCE HTTP/1.1", "T-1"), + ("GET /api/w/workspace-a/tickets/01TARGET HTTP/1.1", "T-2"), + ] { + let (mut stream, _) = listener.accept().unwrap(); + let mut buffer = [0_u8; 8192]; + let len = stream.read(&mut buffer).unwrap(); + let request = String::from_utf8_lossy(&buffer[..len]); + assert!(request.starts_with(expected_path)); + let body = serde_json::json!({"meta": {"resource_key": resource_key}}).to_string(); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", + body.len(), body + ) + .unwrap(); + } + + let (mut stream, _) = listener.accept().unwrap(); + let mut buffer = [0_u8; 8192]; + let len = stream.read(&mut buffer).unwrap(); + let request = String::from_utf8_lossy(&buffer[..len]); + assert!( + request.starts_with("POST /api/w/workspace-a/tickets/01SOURCE/relations HTTP/1.1") + ); + let body = serde_json::to_string(&TicketRelation { + ticket_id: "01SOURCE".to_string(), + kind: TicketRelationKind::DependsOn, + target: "01TARGET".to_string(), + note: None, + author: "worker-internal".to_string(), + at: "2026-08-06T00:00:00Z".to_string(), + }) + .unwrap(); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", + body.len(), + body + ) + .unwrap(); + }); + let backend = WorkspaceHttpTicketBackend::new(Arc::new( + crate::worker::TestWorkspaceHttpClient::new("workspace-a", format!("http://{addr}")), + )); + + let relation = backend + .add_ticket_relation( + TicketIdOrSlug::Id("01SOURCE".to_string()), + NewTicketRelation { + kind: TicketRelationKind::DependsOn, + target: "01TARGET".to_string(), + note: None, + author: None, + }, + ) + .unwrap(); + + server.join().unwrap(); + assert_eq!(relation.ticket_id, "T-1"); + assert_eq!(relation.target, "T-2"); + assert_eq!(relation.author, "workspace"); + } + #[test] fn workspace_http_backend_deletes_exact_ticket_relation() { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let base_url = format!("http://{}", listener.local_addr().unwrap()); let server = thread::spawn(move || { + for (expected_path, resource_key) in [ + ("GET /api/w/workspace-a/tickets/01SOURCE HTTP/1.1", "T-1"), + ("GET /api/w/workspace-a/tickets/01TARGET HTTP/1.1", "T-2"), + ] { + let (mut stream, _) = listener.accept().unwrap(); + let mut buffer = [0_u8; 8192]; + let len = stream.read(&mut buffer).unwrap(); + let request = String::from_utf8_lossy(&buffer[..len]); + assert!(request.starts_with(expected_path)); + let response_body = serde_json::json!({ + "meta": {"resource_key": resource_key} + }) + .to_string(); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", + response_body.len(), + response_body + ) + .unwrap(); + } + let (mut stream, _) = listener.accept().unwrap(); let mut buffer = [0_u8; 8192]; let len = stream.read(&mut buffer).unwrap(); @@ -1787,8 +1996,8 @@ provider = "github" .unwrap(); server.join().unwrap(); - assert_eq!(removed.ticket_id, "01SOURCE"); - assert_eq!(removed.target, "01TARGET"); + assert_eq!(removed.ticket_id, "T-1"); + assert_eq!(removed.target, "T-2"); } #[test] diff --git a/crates/worker/src/feature/builtin/worker_observation.rs b/crates/worker/src/feature/builtin/worker_observation.rs index 2fdbe863..68305b2c 100644 --- a/crates/worker/src/feature/builtin/worker_observation.rs +++ b/crates/worker/src/feature/builtin/worker_observation.rs @@ -1,11 +1,12 @@ use std::sync::Arc; +#[cfg(test)] use agen::Item; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use async_trait::async_trait; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use session_store::collect_state; +use session_store::{LogEntry, collect_state}; use super::manage_worker::{WORKER_CONTROL_SERVICE_ID, WorkerControlService}; use crate::feature::{ @@ -60,7 +61,27 @@ pub struct WorkerObservationSubject { #[derive(Debug, Clone)] pub struct WorkerSessionCapture { pub segment_id: String, - pub items: Vec, + pub entries: Vec>, +} + +impl WorkerSessionCapture { + pub fn from_log_entries( + segment_id: impl Into, + log_entries: &[LogEntry], + ) -> Result { + let segment_id = segment_id.into(); + let state = collect_state(log_entries); + let parsed_segment_id = segment_id.parse().unwrap_or_default(); + let entries = crate::session_history::restore_history_entries( + state.session_id.unwrap_or_default(), + parsed_segment_id, + log_entries, + )?; + Ok(Self { + segment_id, + entries, + }) + } } #[derive(Debug, thiserror::Error)] @@ -161,9 +182,17 @@ impl WorkerObservationProvider for WorkspaceClientWorkerObservationProvider { }) .collect::, _>>()?; let state = collect_state(&entries); + let segment_id = response.segment_id; + let parsed_segment_id = segment_id.parse().unwrap_or_default(); + let typed_entries = crate::session_history::restore_history_entries( + state.session_id.unwrap_or_default(), + parsed_segment_id, + &entries, + ) + .map_err(WorkerObservationError::Unavailable)?; Ok(WorkerSessionCapture { - segment_id: response.segment_id, - items: state.history, + segment_id, + entries: typed_entries, }) } } @@ -392,9 +421,15 @@ impl WorkerObservationProvider for SpawnedSubWorkerObservationProvider { .ok_or(WorkerObservationError::NotFound)?; let entries = record.session.entries(); let state = collect_state(&entries); + let typed_entries = crate::session_history::restore_history_entries( + state.session_id.unwrap_or_default(), + Default::default(), + &entries, + ) + .map_err(WorkerObservationError::Unavailable)?; Ok(WorkerSessionCapture { segment_id: format!("subworker:{name}"), - items: state.history, + entries: typed_entries, }) } } @@ -508,6 +543,7 @@ impl Tool for ViewSessionOverviewTool { .map(|entry| { serde_json::json!({ "entry_ref": entry.id, + "origin": entry.origin, "entry_range": entry.entry_range, "kind": entry.kind.as_str(), "label": entry.label, @@ -547,7 +583,7 @@ impl Tool for SearchSessionEntriesTool { let from = params.from.as_deref().map(parse_entry_ref).transpose()?; let through = params.through.as_deref().map(parse_entry_ref).transpose()?; if let (Some(from), Some(through)) = (&from, &through) { - if from.source_index() > through.source_index() { + if view.source_index_for_ref(from) > view.source_index_for_ref(through) { return Err(ToolError::InvalidArgument( "SearchSessionEntries from must not be after through".to_string(), )); @@ -573,6 +609,7 @@ impl Tool for SearchSessionEntriesTool { .map(|entry| { serde_json::json!({ "entry_ref": entry.id, + "origin": entry.origin, "entry_range": entry.entry_range, "kind": entry.kind.as_str(), "tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()), @@ -628,6 +665,7 @@ impl Tool for ReadSessionEntryTool { .map(|entry| { serde_json::json!({ "entry_ref": entry.id, + "origin": entry.origin, "entry_range": entry.entry_range, "kind": entry.kind.as_str(), "tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()), @@ -661,7 +699,10 @@ async fn latest_view( .capture_worker_session(subject) .await .map_err(tool_error)?; - Ok(SessionCapture::new(capture.segment_id, capture.items)) + Ok(SessionCapture::from_history_entries( + capture.segment_id, + capture.entries, + )) } fn parse_input( @@ -751,9 +792,23 @@ mod tests { if subject != &granted_subject() { return Err(WorkerObservationError::NotFound); } + let entries = self + .captures + .lock() + .unwrap() + .clone() + .into_iter() + .enumerate() + .map(|(index, item)| { + let mut metadata = crate::SessionHistoryMetadata::legacy_unknown(); + metadata.entry_id = + session_store::LoggedSessionHistoryEntryId(format!("fake-{index:08}")); + agen::HistoryEntry::new(item, metadata) + }) + .collect(); Ok(WorkerSessionCapture { segment_id: "segment".to_string(), - items: self.captures.lock().unwrap().clone(), + entries, }) } } @@ -796,7 +851,7 @@ mod tests { let read = read_definition(provider.clone())().1; let hidden = read .execute( - r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"unauthorized"},"entry_ref":"E00000000"}"#, + r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"unauthorized"},"entry_ref":"Efake-00000000"}"#, agen::tool::ToolExecutionContext::direct(), ) .await @@ -810,7 +865,7 @@ mod tests { .push(message("a1", Role::Assistant, "second")); let output = read .execute( - r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"granted"},"entry_ref":"E00000000"}"#, + r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"granted"},"entry_ref":"Efake-00000000"}"#, agen::tool::ToolExecutionContext::direct(), ) .await @@ -819,7 +874,7 @@ mod tests { let output = read .execute( - r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"granted"},"entry_ref":"E00000001"}"#, + r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"granted"},"entry_ref":"Efake-00000001"}"#, agen::tool::ToolExecutionContext::direct(), ) .await diff --git a/crates/worker/src/internal_worker.rs b/crates/worker/src/internal_worker.rs index f47ef183..a8d21e49 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -10,7 +10,7 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex}; use agen::timeline::event::UsageEvent; -use agen::{Engine, llm_client::LlmClient}; +use agen::{Engine, EngineError, llm_client::LlmClient}; use manifest::{Scope, WorkerManifest}; use protocol::{Event, InFlightSnapshot, WorkerStatus}; use session_store::{LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntry}; @@ -55,7 +55,17 @@ pub(crate) struct InternalWorkerSpec { pub input: String, pub cache_key: Option, pub max_turns: Option, - pub engine_configurator: Option>) + Send>>, + pub engine_configurator: Option< + Box< + dyn FnOnce( + &mut Engine< + Box, + agen::state::Mutable, + crate::SessionHistoryMetadata, + >, + ) + Send, + >, + >, pub features: FeatureRegistryBuilder, pub required_tools: &'static [&'static str], pub authority: InternalWorkerAuthority, @@ -124,7 +134,9 @@ where let last_usage = Arc::new(Mutex::new(None::)); let usage_slot = last_usage.clone(); - let mut engine = Engine::new(client).system_prompt(system_prompt); + let mut engine = + Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated(client) + .system_prompt(system_prompt); engine.on_usage(move |usage| { if let Ok(mut slot) = usage_slot.lock() { *slot = Some(usage.clone()); @@ -199,12 +211,28 @@ where on_cancel_sender(worker.engine_mut().cancel_sender()); match worker.run_text(&input).await { - Ok(lifecycle) => Ok(InternalWorkerResult { + Ok(lifecycle @ WorkerRunResult::Finished) + | Ok(lifecycle @ WorkerRunResult::Paused) + | Ok(lifecycle @ WorkerRunResult::RolledBack) => Ok(InternalWorkerResult { usage: last_usage.lock().ok().and_then(|slot| slot.clone()), identity, lifecycle, history_entries: store.entries_count(session_id, segment_id), }), + Ok(WorkerRunResult::LimitReached) => Err(InternalWorkerError { + source: WorkerError::Engine(EngineError::Aborted( + "internal Worker reached its turn limit".to_string(), + )), + usage: last_usage.lock().ok().and_then(|slot| slot.clone()), + identity, + history_entries: store.entries_count(session_id, segment_id), + }), + Ok(WorkerRunResult::Interrupted { message, .. }) => Err(InternalWorkerError { + source: WorkerError::Engine(EngineError::Aborted(message)), + usage: last_usage.lock().ok().and_then(|slot| slot.clone()), + identity, + history_entries: store.entries_count(session_id, segment_id), + }), Err(source) => Err(InternalWorkerError { source, usage: last_usage.lock().ok().and_then(|slot| slot.clone()), @@ -232,6 +260,7 @@ impl Default for InternalWorkerVisibility { pub(crate) enum InternalWorkerSessionStatus { Idle, Running, + Paused, Stopping, Stopped, Failed, @@ -242,9 +271,10 @@ impl InternalWorkerSessionStatus { match self { Self::Idle => 0, Self::Running => 1, - Self::Stopping => 2, - Self::Stopped => 3, - Self::Failed => 4, + Self::Paused => 2, + Self::Stopping => 3, + Self::Stopped => 4, + Self::Failed => 5, } } @@ -252,13 +282,35 @@ impl InternalWorkerSessionStatus { match value { 0 => Self::Idle, 1 => Self::Running, - 2 => Self::Stopping, - 3 => Self::Stopped, + 2 => Self::Paused, + 3 => Self::Stopping, + 4 => Self::Stopped, _ => Self::Failed, } } } +fn classify_internal_turn_result( + result: Result, +) -> (InternalWorkerSessionStatus, Option) { + match result { + Ok(WorkerRunResult::Finished) => (InternalWorkerSessionStatus::Idle, None), + Ok(WorkerRunResult::Paused) => (InternalWorkerSessionStatus::Paused, None), + Ok(WorkerRunResult::LimitReached) => ( + InternalWorkerSessionStatus::Stopped, + Some("internal Worker reached its turn limit".to_string()), + ), + Ok(WorkerRunResult::Interrupted { message, .. }) => { + (InternalWorkerSessionStatus::Stopped, Some(message)) + } + Ok(WorkerRunResult::RolledBack) => ( + InternalWorkerSessionStatus::Stopped, + Some("internal Worker run was cancelled before AI output".to_string()), + ), + Err(error) => (InternalWorkerSessionStatus::Failed, Some(error.to_string())), + } +} + #[derive(Debug, thiserror::Error)] pub(crate) enum InternalWorkerSessionError { #[error("failed to build internal Worker session: {message}")] @@ -353,10 +405,11 @@ impl InternalWorkerSessionHandle { entries, status: match self.status() { InternalWorkerSessionStatus::Running => WorkerStatus::Running, + InternalWorkerSessionStatus::Paused => WorkerStatus::Paused, InternalWorkerSessionStatus::Idle => WorkerStatus::Idle, InternalWorkerSessionStatus::Stopping | InternalWorkerSessionStatus::Stopped - | InternalWorkerSessionStatus::Failed => WorkerStatus::Paused, + | InternalWorkerSessionStatus::Failed => WorkerStatus::Stopped, }, error: self.last_error.lock().unwrap().clone(), in_flight, @@ -388,6 +441,7 @@ impl InternalWorkerSessionHandle { .map_err( |current| match InternalWorkerSessionStatus::decode(current) { InternalWorkerSessionStatus::Running + | InternalWorkerSessionStatus::Paused | InternalWorkerSessionStatus::Stopping => InternalWorkerSessionError::Busy, InternalWorkerSessionStatus::Stopped | InternalWorkerSessionStatus::Failed => { InternalWorkerSessionError::Stopped @@ -494,7 +548,9 @@ pub(crate) async fn spawn_internal_worker_session( let last_usage = Arc::new(Mutex::new(None::)); let usage_slot = last_usage.clone(); - let mut engine = Engine::new(client).system_prompt(system_prompt); + let mut engine = + Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated(client) + .system_prompt(system_prompt); engine.on_usage(move |usage| { if let Ok(mut slot) = usage_slot.lock() { *slot = Some(usage.clone()); @@ -591,7 +647,9 @@ pub(crate) fn prepare_internal_worker_from_spec( manifest.compaction = None; manifest.memory = None; - let mut engine = Engine::new(client).system_prompt(system_prompt); + let mut engine = + Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated(client) + .system_prompt(system_prompt); engine.set_cache_key(cache_key); engine.set_max_turns(max_turns); if let Some(configure) = engine_configurator { @@ -733,13 +791,7 @@ pub(crate) async fn prepare_internal_worker_session( loop { tokio::select! { result = &mut run => { - let (turn_status, error) = match result { - Ok(_) => (InternalWorkerSessionStatus::Idle, None), - Err(error) => ( - InternalWorkerSessionStatus::Failed, - Some(error.to_string()), - ), - }; + let (turn_status, error) = classify_internal_turn_result(result); actor_in_flight.clear(); status.store(turn_status.encode(), std::sync::atomic::Ordering::Release); if let Some(message) = error { @@ -748,11 +800,20 @@ pub(crate) async fn prepare_internal_worker_session( code: protocol::ErrorCode::Internal, message, }); - } else { - let _ = event_tx.send(Event::Status { - status: WorkerStatus::Idle, - }); } + let protocol_status = match turn_status { + InternalWorkerSessionStatus::Idle => WorkerStatus::Idle, + InternalWorkerSessionStatus::Paused => WorkerStatus::Paused, + InternalWorkerSessionStatus::Stopped + | InternalWorkerSessionStatus::Failed => WorkerStatus::Stopped, + InternalWorkerSessionStatus::Running + | InternalWorkerSessionStatus::Stopping => { + unreachable!("run completion cannot remain active") + } + }; + let _ = event_tx.send(Event::Status { + status: protocol_status, + }); if let Some(callback) = &on_turn_end { callback(turn_status); } @@ -766,7 +827,7 @@ pub(crate) async fn prepare_internal_worker_session( let _ = (&mut run).await; actor_in_flight.clear(); status.store(InternalWorkerSessionStatus::Stopped.encode(), std::sync::atomic::Ordering::Release); - let _ = event_tx.send(Event::Status { status: WorkerStatus::Paused }); + let _ = event_tx.send(Event::Status { status: WorkerStatus::Stopped }); let _ = event_tx.send(Event::Shutdown); state_changed.notify_waiters(); let _ = done.send(()); @@ -792,7 +853,7 @@ pub(crate) async fn prepare_internal_worker_session( std::sync::atomic::Ordering::Release, ); let _ = event_tx.send(Event::Status { - status: WorkerStatus::Paused, + status: WorkerStatus::Stopped, }); let _ = event_tx.send(Event::Shutdown); state_changed.notify_waiters(); @@ -1102,6 +1163,26 @@ mod tests { } } + #[derive(Clone)] + struct FailingClient; + + #[async_trait] + impl LlmClient for FailingClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + + async fn stream( + &self, + _request: Request, + ) -> Result> + Send>>, ClientError> + { + Err(ClientError::Config( + "intentional internal failure".to_string(), + )) + } + } + #[derive(Clone)] struct CancelBeforeAiClient { calls: Arc, @@ -1215,6 +1296,77 @@ permission = "write" assert_eq!(result.identity.kind, "test"); } + #[test] + fn internal_turn_result_mapping_is_exhaustive() { + let cases = [ + ( + WorkerRunResult::Finished, + InternalWorkerSessionStatus::Idle, + false, + ), + ( + WorkerRunResult::Paused, + InternalWorkerSessionStatus::Paused, + false, + ), + ( + WorkerRunResult::LimitReached, + InternalWorkerSessionStatus::Stopped, + true, + ), + ( + WorkerRunResult::Interrupted { + code: protocol::ErrorCode::Internal, + message: "cancelled".to_string(), + }, + InternalWorkerSessionStatus::Stopped, + true, + ), + ( + WorkerRunResult::RolledBack, + InternalWorkerSessionStatus::Stopped, + true, + ), + ]; + + for (result, expected_status, expects_error) in cases { + let (status, error) = classify_internal_turn_result(Ok(result)); + assert_eq!(status, expected_status); + assert_eq!(error.is_some(), expects_error); + } + + let (status, error) = classify_internal_turn_result(Err(WorkerError::Engine( + EngineError::Aborted("fatal".to_string()), + ))); + assert_eq!(status, InternalWorkerSessionStatus::Failed); + assert!(error.is_some_and(|message| message.contains("fatal"))); + } + + #[tokio::test] + async fn fatal_internal_run_transitions_to_stopped_protocol_status() { + let calls = Arc::new(AtomicUsize::new(0)); + let mut internal_spec = spec(calls, &[]); + internal_spec.client = Box::new(FailingClient); + + let handle = spawn_internal_worker_session(internal_spec) + .await + .expect("spawn failing Internal Worker session"); + assert_eq!( + handle.wait_until_idle().await, + InternalWorkerSessionStatus::Stopped + ); + assert_eq!(handle.status(), InternalWorkerSessionStatus::Stopped); + assert_eq!(handle.protocol_snapshot().status, WorkerStatus::Stopped); + assert!( + handle + .last_error + .lock() + .unwrap() + .as_ref() + .is_some_and(|message| message.contains("intentional internal failure")) + ); + } + #[tokio::test] async fn session_accepts_follow_up_turns_and_stops_without_runtime_registration() { let calls = Arc::new(AtomicUsize::new(0)); diff --git a/crates/worker/src/interrupt_prep.rs b/crates/worker/src/interrupt_prep.rs index 6e359828..b02a7284 100644 --- a/crates/worker/src/interrupt_prep.rs +++ b/crates/worker/src/interrupt_prep.rs @@ -13,7 +13,7 @@ #[cfg(test)] use crate::prompt::catalog::PromptCatalog; -use agen::Item; +use agen::{Item, ToolResultDisposition}; /// Build synthetic `Item::ToolResult` items for every unanswered /// `Item::ToolCall` in `history`, preserving order. @@ -28,7 +28,16 @@ pub(crate) fn orphan_tool_result_closures(history: &[Item], summary: &str) -> Ve for item in history { if let Item::ToolCall { call_id, .. } = item { if !answered.contains(call_id.as_str()) { - out.push(Item::tool_result(call_id.clone(), summary)); + out.push(Item::tool_result_item_with_disposition_and_attachments( + call_id.clone(), + summary, + Some( + "Execution ended before completion could be confirmed. Completion and side effects are unknown." + .to_string(), + ), + ToolResultDisposition::OutcomeUnknown, + Vec::new(), + )); } } } @@ -77,10 +86,12 @@ mod tests { Item::ToolResult { call_id, summary: got, + disposition, .. } => { assert_eq!(call_id, "c1"); assert_eq!(got, &summary); + assert_eq!(*disposition, ToolResultDisposition::OutcomeUnknown); } other => panic!("expected ToolResult, got {other:?}"), } diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index a99fc7a3..a6283f25 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -8,6 +8,7 @@ //! decisions (continue / skip / abort / pause). use std::borrow::Cow; +use std::collections::VecDeque; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; @@ -33,7 +34,9 @@ use crate::hook::{ }; use crate::ipc::notify_buffer::{NotifyBuffer, build_system_item_with_provenance}; use crate::prompt::catalog::PromptCatalog; +use crate::session_history::SessionHistoryMetadata; use crate::worker::SystemItemCommitter; +use agen::HistoryEntry; use agen::token_counter::total_tokens; /// Maximum number of bytes copied into `TurnEndInfo::final_text_preview`. @@ -73,6 +76,7 @@ pub(crate) struct WorkerInterceptor { /// worker. `None` in tests / `Worker::new` paths where no writer is /// attached. log_writer: Option>, + pending_committed_history: Arc>>>, /// Next turn index assigned by `on_prompt_submit`. next_turn_index: AtomicUsize, /// Tool calls observed in the current turn (reset on each new prompt). @@ -80,6 +84,7 @@ pub(crate) struct WorkerInterceptor { } impl WorkerInterceptor { + #[cfg(test)] pub(crate) fn new( registry: Arc, compact_state: Option>, @@ -88,6 +93,28 @@ impl WorkerInterceptor { pending_attachments: Arc>>, prompts: Arc>, log_writer: Option>, + ) -> Self { + Self::new_with_history_queue( + registry, + compact_state, + usage_history, + pending_notifies, + pending_attachments, + prompts, + log_writer, + Arc::new(Mutex::new(VecDeque::new())), + ) + } + + pub(crate) fn new_with_history_queue( + registry: Arc, + compact_state: Option>, + usage_history: Option>>>, + pending_notifies: NotifyBuffer, + pending_attachments: Arc>>, + prompts: Arc>, + log_writer: Option>, + pending_committed_history: Arc>>>, ) -> Self { Self { registry, @@ -99,6 +126,7 @@ impl WorkerInterceptor { prompts, prompt_workspace_id: None, log_writer, + pending_committed_history, next_turn_index: AtomicUsize::new(0), tool_calls_this_turn: AtomicUsize::new(0), } @@ -125,7 +153,11 @@ impl WorkerInterceptor { return Ok(()); }; for item in items { - writer.commit_system_item(item.clone())?; + let entry = writer.commit_system_item(item.clone())?; + self.pending_committed_history + .lock() + .expect("pending committed history poisoned") + .push_back(entry); } Ok(()) } @@ -507,7 +539,12 @@ mod tests { &self, entry: session_store::LogEntry, ) -> Result<(), session_store::StoreError> { - if let session_store::LogEntry::SystemItem { item, .. } = entry { + let item = match entry { + session_store::LogEntry::SystemItem { item, .. } => Some(item), + session_store::LogEntry::AnnotatedSystemItem { entry, .. } => Some(entry.item), + _ => None, + }; + if let Some(item) = item { self.committed .lock() .expect("committed system-item list poisoned") diff --git a/crates/worker/src/ipc/protocol_session.rs b/crates/worker/src/ipc/protocol_session.rs index 873d94d0..9b749126 100644 --- a/crates/worker/src/ipc/protocol_session.rs +++ b/crates/worker/src/ipc/protocol_session.rs @@ -29,15 +29,21 @@ pub fn subscribe_worker_protocol_session(handle: &WorkerHandle) -> WorkerProtoco pub fn live_log_entry_event(entry: LogEntry) -> Option { match entry { - LogEntry::SegmentStart { .. } => { + entry @ (LogEntry::SegmentStart { .. } | LogEntry::AnnotatedSegmentStart { .. }) => { let value = serde_json::to_value(&entry).expect("LogEntry is Serialize"); Some(Event::SegmentRotated { entry: value }) } - LogEntry::UserInput { segments, .. } => Some(Event::UserMessage { segments }), + LogEntry::UserInput { segments, .. } | LogEntry::AnnotatedUserInput { segments, .. } => { + Some(Event::UserMessage { segments }) + } LogEntry::SystemItem { item, .. } => { let value = serde_json::to_value(&item).expect("SystemItem is Serialize"); Some(Event::SystemItem { item: value }) } + LogEntry::AnnotatedSystemItem { entry, .. } => { + let value = serde_json::to_value(&entry.item).expect("SystemItem is Serialize"); + Some(Event::SystemItem { item: value }) + } LogEntry::Invoke { trigger, .. } => Some(Event::InvokeStart { kind: trigger }), other => { // `SegmentLogSink::is_live_relevant` keeps non-live-relevant diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 76ced8ce..8de12fbd 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -12,6 +12,7 @@ pub mod prompt; pub mod runtime; pub mod segment_log_sink; mod session_capture; +mod session_history; pub mod shared_state; mod shutdown_after_idle; pub mod skill; @@ -33,14 +34,19 @@ pub use manifest::{ }; pub use model_client::{ProviderError, build_client}; pub use prompt::catalog::{ - CatalogError, EffectivePromptCatalog, PromptCatalog, WorkerPrompt, WorkspacePromptProjection, - prompt_schema_source, + CatalogError, EffectivePromptCatalog, OrchestratorQueueAttentionContext, + OrchestratorQueueAttentionPrompt, OrchestratorQueueAttentionTicket, PromptCatalog, + WorkerPrompt, WorkspacePromptProjection, prompt_schema_source, }; pub use prompt::source::PromptCatalogSource; pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTemplate}; pub use protocol::{ErrorCode, Event, Method, TurnResult, WorkerStatus}; pub use runtime::dir::RuntimeDir; pub use segment_log_sink::SegmentLogSink; +pub use session_history::{ + SessionHistoryDerivation, SessionHistoryEntryId, SessionHistoryMetadata, + WorkerHistoryProvenance, WorkerSubjectSnapshot, +}; pub use shared_state::WorkerSharedState; pub use worker::{ LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError, diff --git a/crates/worker/src/permission.rs b/crates/worker/src/permission.rs index 2b1f28c8..22b5c831 100644 --- a/crates/worker/src/permission.rs +++ b/crates/worker/src/permission.rs @@ -34,7 +34,7 @@ impl PermissionHook { } } -impl Worker { +impl Worker { pub(crate) fn apply_permissions_from_manifest(&mut self) { let Some(permissions) = self.manifest().permissions.clone() else { return; diff --git a/crates/worker/src/prompt/catalog.rs b/crates/worker/src/prompt/catalog.rs index a02d7871..87986cff 100644 --- a/crates/worker/src/prompt/catalog.rs +++ b/crates/worker/src/prompt/catalog.rs @@ -141,8 +141,93 @@ impl WorkerPrompt { ]; } +/// Model-visible queued Ticket projection shared by Server and TUI backlog attention paths. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct OrchestratorQueueAttentionTicket { + resource_key: String, + title: String, +} + +impl OrchestratorQueueAttentionTicket { + pub fn new( + resource_key: impl Into, + title: impl Into, + ) -> Result { + let resource_key = resource_key.into(); + if !is_ticket_resource_key(&resource_key) { + return Err(CatalogError::InvalidQueueAttentionResourceKey); + } + Ok(Self { + resource_key, + title: bounded_queue_attention_text(&title.into(), 240), + }) + } +} + +/// Shared model-visible context for every Orchestrator backlog attention renderer. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct OrchestratorQueueAttentionContext { + tickets: Vec, + separator: &'static str, + omitted_ticket_count: usize, +} + +impl OrchestratorQueueAttentionContext { + pub const MAX_TICKETS: usize = 20; + + pub fn new(tickets: Vec) -> Self { + let omitted_ticket_count = tickets.len().saturating_sub(Self::MAX_TICKETS); + Self { + tickets: tickets.into_iter().take(Self::MAX_TICKETS).collect(), + separator: "—", + omitted_ticket_count, + } + } +} + +/// Prompt-catalog entries that must share the same backlog-attention body contract. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OrchestratorQueueAttentionPrompt { + Server, + Tui, +} + +impl OrchestratorQueueAttentionPrompt { + fn key(self) -> &'static str { + match self { + Self::Server => "internal.workspace_orchestrator_queue_attention", + Self::Tui => "panel.orchestrator_idle_queue_notice", + } + } +} + +fn is_ticket_resource_key(input: &str) -> bool { + input.len() <= 32 + && input.strip_prefix("T-").is_some_and(|suffix| { + !suffix.is_empty() && suffix.bytes().all(|byte| byte.is_ascii_digit()) + }) +} + +fn bounded_queue_attention_text(input: &str, max_chars: usize) -> String { + let mut output = String::new(); + for (index, character) in input.chars().enumerate() { + if index == max_chars { + output.push('…'); + break; + } + output.push(if character.is_control() { + ' ' + } else { + character + }); + } + output +} + #[derive(Debug, Error)] pub enum CatalogError { + #[error("queued Ticket resource key is missing or invalid")] + InvalidQueueAttentionResourceKey, #[error("failed to build builtin Prompt source tree: {0}")] BuiltinTree(String), #[error("failed to evaluate builtin Prompt source tree: {0}")] @@ -319,6 +404,14 @@ impl PromptCatalog { self.render_name(key, Value::from_serialize(context)) } + pub fn orchestrator_queue_attention( + &self, + prompt: OrchestratorQueueAttentionPrompt, + context: &OrchestratorQueueAttentionContext, + ) -> Result { + self.render_serializable(prompt.key(), context) + } + pub fn render_name(&self, key: &str, ctx: Value) -> Result { let template = self .env @@ -653,6 +746,62 @@ mod tests { assert!(reviewer.contains("target-only movement does not invalidate approval")); } + #[test] + fn queue_attention_prompts_share_sanitized_contract_and_true_truncation() { + let catalog = PromptCatalog::builtins_only().unwrap(); + let tickets = (1..=OrchestratorQueueAttentionContext::MAX_TICKETS + 1) + .map(|index| { + OrchestratorQueueAttentionTicket::new( + format!("T-{index}"), + format!("Ticket {index}\nwith control\u{7}"), + ) + .unwrap() + }) + .collect(); + let context = OrchestratorQueueAttentionContext::new(tickets); + let server = catalog + .orchestrator_queue_attention(OrchestratorQueueAttentionPrompt::Server, &context) + .unwrap(); + let tui = catalog + .orchestrator_queue_attention(OrchestratorQueueAttentionPrompt::Tui, &context) + .unwrap(); + + assert_eq!(server, tui); + assert!(server.starts_with("Queued Tickets require attention:")); + assert!(server.contains("- T-1 — Ticket 1 with control ")); + assert!(!server.contains("T-21")); + assert!(server.contains("were omitted from this notice: 1")); + assert!(server.contains("Re-query current Ticket authority")); + assert!(server.contains("Reread the current Ticket state before acting")); + for secret in [ + "workspace_id", + "Workspace:", + "runtime_id", + "worker_id", + "bounded", + ] { + assert!(!server.contains(secret), "leaked {secret}: {server}"); + } + } + + #[test] + fn queue_attention_prompt_omits_truncation_text_for_complete_list() { + let catalog = PromptCatalog::builtins_only().unwrap(); + let context = OrchestratorQueueAttentionContext::new(vec![ + OrchestratorQueueAttentionTicket::new("T-541", "Attention contract").unwrap(), + ]); + let rendered = catalog + .orchestrator_queue_attention(OrchestratorQueueAttentionPrompt::Server, &context) + .unwrap(); + + assert!(rendered.contains("- T-541 — Attention contract")); + assert!(!rendered.contains("omitted")); + assert!(matches!( + OrchestratorQueueAttentionTicket::new("opaque-id", "must fail"), + Err(CatalogError::InvalidQueueAttentionResourceKey) + )); + } + #[test] fn graph_rejects_dynamic_legacy_missing_and_cycles() { let invalid = BTreeMap::from([ diff --git a/crates/worker/src/segment_log_sink.rs b/crates/worker/src/segment_log_sink.rs index 8c4a1f3f..77ee51eb 100644 --- a/crates/worker/src/segment_log_sink.rs +++ b/crates/worker/src/segment_log_sink.rs @@ -121,8 +121,11 @@ impl SegmentLogSink { matches!( entry, LogEntry::SegmentStart { .. } + | LogEntry::AnnotatedSegmentStart { .. } | LogEntry::UserInput { .. } + | LogEntry::AnnotatedUserInput { .. } | LogEntry::SystemItem { .. } + | LogEntry::AnnotatedSystemItem { .. } | LogEntry::Invoke { .. } ) } diff --git a/crates/worker/src/session_capture.rs b/crates/worker/src/session_capture.rs index d3ae8ec6..1886e4e8 100644 --- a/crates/worker/src/session_capture.rs +++ b/crates/worker/src/session_capture.rs @@ -6,7 +6,8 @@ use std::sync::Arc; -use agen::{Item, Role}; +use crate::session_history::{SessionHistoryMetadata, WorkerHistoryProvenance}; +use agen::{HistoryEntry, Item, Role}; use serde::{Deserialize, Serialize}; const DEFAULT_SEARCH_LIMIT: usize = 20; @@ -21,14 +22,21 @@ const OVERVIEW_ANCHOR_STRIDE: usize = 8; pub(crate) struct SessionEntryRef(String); impl SessionEntryRef { - pub(crate) fn new(source_index: usize) -> Self { - Self(format!("E{source_index:08}")) + pub(crate) fn from_history_entry_id(entry_id: &crate::SessionHistoryEntryId) -> Self { + Self(format!("E{}", entry_id.0)) } pub(crate) fn parse(value: &str) -> Option { - let reference = Self(value.to_string()); - reference.source_index()?; - Some(reference) + let suffix = value.strip_prefix('E')?; + if suffix.is_empty() + || suffix.len() > 64 + || !suffix + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-') + { + return None; + } + Some(Self(value.to_string())) } pub(crate) fn as_str(&self) -> &str { @@ -97,6 +105,7 @@ impl ToolPart { #[derive(Debug, Clone)] pub(crate) struct OverviewItem { pub id: SessionEntryRef, + pub origin: WorkerHistoryProvenance, pub entry_range: [u64; 2], pub kind: ReferenceKind, pub label: String, @@ -107,6 +116,7 @@ pub(crate) struct OverviewItem { #[derive(Debug, Clone)] pub(crate) struct ReferenceEntry { pub id: SessionEntryRef, + pub origin: WorkerHistoryProvenance, pub entry_range: [u64; 2], pub kind: ReferenceKind, pub tool_part: Option, @@ -132,6 +142,7 @@ pub(crate) struct SearchOptions { #[derive(Debug, Clone)] pub(crate) struct SearchHit { pub id: SessionEntryRef, + pub origin: WorkerHistoryProvenance, pub kind: ReferenceKind, pub tool_part: Option, pub tool_name: Option, @@ -177,6 +188,7 @@ impl Default for ReadOptions { #[derive(Debug, Clone)] pub(crate) struct ReadEntry { pub id: SessionEntryRef, + pub origin: WorkerHistoryProvenance, pub kind: ReferenceKind, pub tool_part: Option, pub tool_name: Option, @@ -195,6 +207,7 @@ pub(crate) struct ReadResult { pub(crate) struct SessionEntryEvidence { pub segment_id: String, pub entry_ref: SessionEntryRef, + pub origin: WorkerHistoryProvenance, pub entry_range: [u64; 2], pub kind: ReferenceKind, pub tool_part: Option, @@ -206,26 +219,42 @@ pub(crate) struct SessionEntryEvidence { #[derive(Debug, Clone)] pub(crate) struct SessionCapture { segment_id: String, - items: Arc>, + entries: Arc>>, overview: Vec, index: Vec, } impl SessionCapture { pub(crate) fn new(segment_id: impl Into, items: Vec) -> Self { + let entries = items + .into_iter() + .enumerate() + .map(|(index, item)| { + let mut metadata = SessionHistoryMetadata::legacy_unknown(); + metadata.entry_id = + session_store::LoggedSessionHistoryEntryId(format!("{index:08}")); + HistoryEntry::new(item, metadata) + }) + .collect(); + Self::from_history_entries(segment_id, entries) + } + + pub(crate) fn from_history_entries( + segment_id: impl Into, + entries: Vec>, + ) -> Self { let segment_id = segment_id.into(); - let items = Arc::new(items); + let entries = Arc::new(entries); let mut overview = Vec::new(); let mut index = Vec::new(); - for (idx, item) in items.iter().enumerate() { + for (idx, entry) in entries.iter().enumerate() { + let item = &entry.item; let entry_range = [idx as u64, idx as u64]; match item { Item::Message { role, content, .. } => { - let kind = match role { - Role::User => ReferenceKind::User, - Role::Assistant => ReferenceKind::Assistant, - Role::System => continue, + let Some(kind) = message_reference_kind(&entry.annotation.origin, role) else { + continue; }; let text = content .iter() @@ -234,9 +263,10 @@ impl SessionCapture { .join(""); let label = format!("{} message", kind.as_str()); let summary = truncate_chars(&text, 240); - let id = SessionEntryRef::new(idx); + let id = SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id); index.push(ReferenceEntry { id: id.clone(), + origin: entry.annotation.origin.clone(), entry_range, kind, tool_part: None, @@ -248,6 +278,7 @@ impl SessionCapture { if matches!(kind, ReferenceKind::User | ReferenceKind::Assistant) { overview.push(OverviewItem { id: id.clone(), + origin: entry.annotation.origin.clone(), entry_range, kind, label, @@ -261,7 +292,8 @@ impl SessionCapture { } => { let text = format!("{name}\n{arguments}"); index.push(ReferenceEntry { - id: SessionEntryRef::new(idx), + id: SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id), + origin: entry.annotation.origin.clone(), entry_range, kind: ReferenceKind::Tool, tool_part: Some(ToolPart::Input), @@ -287,7 +319,8 @@ impl SessionCapture { content.as_deref().unwrap_or_default(), ); index.push(ReferenceEntry { - id: SessionEntryRef::new(idx), + id: SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id), + origin: entry.annotation.origin.clone(), entry_range, kind: ReferenceKind::Tool, tool_part: Some(ToolPart::Output), @@ -327,7 +360,7 @@ impl SessionCapture { Self { segment_id, - items, + entries, overview, index, } @@ -337,6 +370,14 @@ impl SessionCapture { &self.overview } + pub(crate) fn source_index_for_ref(&self, reference: &SessionEntryRef) -> Option { + self.index + .iter() + .find(|entry| entry.id == *reference) + .map(|entry| entry.entry_range[0]) + .or_else(|| reference.source_index()) + } + pub(crate) fn search(&self, options: &SearchOptions) -> Vec { let query = options.query.trim().to_lowercase(); let limit = options @@ -347,12 +388,12 @@ impl SessionCapture { let min_entry_index = options .from .as_ref() - .and_then(SessionEntryRef::source_index) + .and_then(|reference| self.source_index_for_ref(reference)) .unwrap_or_else(|| options.min_entry_index.unwrap_or(0)); let max_entry_index = options .through .as_ref() - .and_then(SessionEntryRef::source_index) + .and_then(|reference| self.source_index_for_ref(reference)) .unwrap_or(u64::MAX); let mut skipped = 0usize; let mut hits = Vec::new(); @@ -391,6 +432,7 @@ impl SessionCapture { } hits.push(SearchHit { id: entry.id.clone(), + origin: entry.origin.clone(), kind: entry.kind, tool_part: entry.tool_part, tool_name: entry.tool_name.clone(), @@ -442,13 +484,18 @@ impl SessionCapture { } } } - let Some(item) = self.items.get(entry.entry_range[0] as usize) else { + let Some(item) = self + .entries + .get(entry.entry_range[0] as usize) + .map(|entry| &entry.item) + else { continue; }; let text = render_item(item, entry, options.detail, max_bytes.saturating_sub(bytes)); bytes = bytes.saturating_add(text.len()); entries.push(ReadEntry { id: entry.id.clone(), + origin: entry.origin.clone(), kind: entry.kind, tool_part: entry.tool_part, tool_name: entry.tool_name.clone(), @@ -485,6 +532,7 @@ impl SessionCapture { Some(SessionEntryEvidence { segment_id: self.segment_id.clone(), entry_ref: entry.id.clone(), + origin: entry.origin.clone(), entry_range: entry.entry_range, kind: entry.kind, tool_part: entry.tool_part, @@ -495,6 +543,28 @@ impl SessionCapture { } } +fn message_reference_kind( + origin: &WorkerHistoryProvenance, + provider_role: &Role, +) -> Option { + match origin { + WorkerHistoryProvenance::HumanInput { .. } + | WorkerHistoryProvenance::WorkerInput { .. } => Some(ReferenceKind::User), + WorkerHistoryProvenance::ModelOutput { .. } => Some(ReferenceKind::Assistant), + WorkerHistoryProvenance::ToolOutput { .. } => Some(ReferenceKind::Tool), + WorkerHistoryProvenance::LegacyUnknown => match provider_role { + Role::User => Some(ReferenceKind::User), + Role::Assistant => Some(ReferenceKind::Assistant), + Role::System => None, + }, + // Flow/backend/system content remains out of the observation surface + // even when represented with a provider user/system role. + WorkerHistoryProvenance::FlowInstruction { .. } + | WorkerHistoryProvenance::BackendInstruction { .. } + | WorkerHistoryProvenance::DerivedSummary => None, + } +} + fn render_item( item: &Item, entry: &ReferenceEntry, @@ -563,6 +633,60 @@ fn truncate_chars(text: &str, max_chars: usize) -> String { mod tests { use super::*; + #[test] + fn flow_user_role_is_excluded_while_explicit_human_origin_remains_evidence() { + let entries = vec![ + crate::session_history::history_entry( + Item::user_message("trusted flow instruction"), + WorkerHistoryProvenance::FlowInstruction { + selector: "builtin:coder-review".into(), + definition_id: "coder-review".into(), + definition_revision: 3, + instance_id: "instance".into(), + state_id: "implement".into(), + }, + ), + crate::session_history::history_entry( + Item::user_message("remember my preference"), + WorkerHistoryProvenance::HumanInput { + account_id: "account-1".into(), + }, + ), + ]; + let capture = SessionCapture::from_history_entries("segment", entries); + let overview = capture.overview(); + assert_eq!(overview.len(), 1); + assert!(matches!( + overview[0].origin, + WorkerHistoryProvenance::HumanInput { .. } + )); + let evidence = capture.evidence_for(overview[0].id.as_str()).unwrap(); + assert!(evidence.excerpt.ends_with("remember my preference")); + assert!(matches!( + evidence.origin, + WorkerHistoryProvenance::HumanInput { .. } + )); + } + + #[test] + fn stable_logical_ref_survives_retention_and_restore_projection() { + let retained = crate::session_history::history_entry( + Item::assistant_message("retained"), + WorkerHistoryProvenance::ModelOutput { + worker: crate::session_history::worker_subject(Default::default()), + }, + ); + let expected_ref = SessionEntryRef::from_history_entry_id(&retained.annotation.entry_id); + let before = SessionCapture::from_history_entries("old", vec![retained.clone()]); + let after = SessionCapture::from_history_entries("new", vec![retained]); + assert_eq!(before.overview()[0].id, expected_ref); + assert_eq!(after.overview()[0].id, expected_ref); + assert_eq!( + after.evidence_for(expected_ref.as_str()).unwrap().entry_ref, + expected_ref + ); + } + #[test] fn overview_contains_user_and_assistant_only() { let view = SessionCapture::new( diff --git a/crates/worker/src/session_history.rs b/crates/worker/src/session_history.rs new file mode 100644 index 00000000..3772d554 --- /dev/null +++ b/crates/worker/src/session_history.rs @@ -0,0 +1,219 @@ +//! Restore-authoritative metadata for model-visible Worker history. +//! +//! Agen transports this annotation without interpreting it. Session Log v2 +//! stores each item and metadata in one typed record; legacy records are +//! retained only as explicit `LegacyUnknown` entries. + +use agen::{HistoryEntry, Item}; +use protocol::Segment; +use session_store::{ + LogEntry, LoggedHistoryDerivation, LoggedHistoryEntry, LoggedSessionHistoryEntryId, + LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, LoggedWorkerSubject, SegmentId, + SessionId, +}; + +pub type SessionHistoryEntryId = LoggedSessionHistoryEntryId; +pub type SessionHistoryMetadata = LoggedSessionHistoryMetadata; +pub type WorkerHistoryProvenance = LoggedSessionHistoryOrigin; +pub type SessionHistoryDerivation = LoggedHistoryDerivation; +pub type WorkerSubjectSnapshot = LoggedWorkerSubject; + +pub(crate) fn worker_subject(session_id: SessionId) -> WorkerSubjectSnapshot { + WorkerSubjectSnapshot { + workspace_id: None, + runtime_id: None, + worker_id: session_id.to_string(), + } +} + +pub(crate) fn metadata( + origin: WorkerHistoryProvenance, + derivation: Option, +) -> SessionHistoryMetadata { + SessionHistoryMetadata { + entry_id: SessionHistoryEntryId::new(), + origin, + derivation, + } +} + +pub(crate) fn history_entry( + item: Item, + origin: WorkerHistoryProvenance, +) -> HistoryEntry { + HistoryEntry::new(item, metadata(origin, None)) +} + +pub(crate) fn to_logged_history_entry( + entry: &HistoryEntry, +) -> LoggedHistoryEntry { + LoggedHistoryEntry { + item: entry.item.clone().into(), + metadata: entry.annotation.clone(), + } +} + +fn legacy_entry(item: Item) -> HistoryEntry { + HistoryEntry::new(item, SessionHistoryMetadata::legacy_unknown()) +} + +fn from_logged(entry: &LoggedHistoryEntry) -> HistoryEntry { + HistoryEntry::new(Item::from(entry.item.clone()), entry.metadata.clone()) +} + +/// Rebuild typed Worker history directly from the append-only Session Log. +/// Missing legacy metadata is never inferred from role or plaintext. +pub(crate) fn restore_history_entries( + _session_id: SessionId, + _segment_id: SegmentId, + entries: &[LogEntry], +) -> Result>, String> { + let mut history = Vec::new(); + for entry in entries { + match entry { + LogEntry::AnnotatedSegmentStart { history: seed, .. } => { + history = seed.iter().map(from_logged).collect(); + } + LogEntry::SegmentStart { history: seed, .. } => { + history = seed + .iter() + .cloned() + .map(Item::from) + .map(legacy_entry) + .collect(); + } + LogEntry::AnnotatedUserInput { history: input, .. } => { + history.extend(input.iter().map(from_logged)) + } + LogEntry::UserInput { segments, .. } => history.push(legacy_entry(Item::user_message( + Segment::flatten_to_text(segments), + ))), + LogEntry::AnnotatedAssistantItem { entry, .. } + | LogEntry::AnnotatedToolResult { entry, .. } => history.push(from_logged(entry)), + LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => { + history.push(legacy_entry(Item::from(item.clone()))); + } + LogEntry::AnnotatedSystemItem { entry, .. } => history.push(HistoryEntry::new( + entry.item.to_history_item(), + entry.metadata.clone(), + )), + LogEntry::SystemItem { item, .. } => { + history.push(legacy_entry(item.to_history_item())); + } + _ => {} + } + } + Ok(history) +} + +#[cfg(test)] +mod tests { + use super::*; + use agen::llm_client::RequestConfig; + use session_store::LogEntry; + + #[test] + fn legacy_user_role_is_not_inferred_as_human_authority() { + let entries = vec![LogEntry::UserInput { + ts: 1, + segments: vec![Segment::text("legacy")], + extensions: Vec::new(), + }]; + let restored = + restore_history_entries(SessionId::now_v7(), SegmentId::now_v7(), &entries).unwrap(); + assert!(matches!( + restored[0].annotation.origin, + WorkerHistoryProvenance::LegacyUnknown + )); + } + + #[test] + fn typed_flow_and_unknown_caller_input_round_trip_without_role_inference() { + let session_id = SessionId::now_v7(); + let projected = vec![ + history_entry( + Item::user_message("flow instructions"), + WorkerHistoryProvenance::FlowInstruction { + selector: "builtin:coder-review".to_string(), + definition_id: "coder-review".to_string(), + definition_revision: 7, + instance_id: "flow-instance".to_string(), + state_id: "implement".to_string(), + }, + ), + history_entry( + Item::user_message("implement"), + WorkerHistoryProvenance::LegacyUnknown, + ), + ]; + let entries = vec![ + LogEntry::AnnotatedSegmentStart { + ts: 0, + session_id, + system_prompt: None, + config: RequestConfig::default(), + history: Vec::new(), + forked_from: None, + compacted_from: None, + }, + LogEntry::AnnotatedUserInput { + ts: 1, + segments: vec![ + Segment::Flow { + selector: "builtin:coder-review".to_string(), + }, + Segment::text("implement"), + ], + extensions: Vec::new(), + history: projected.iter().map(to_logged_history_entry).collect(), + }, + ]; + let restored = restore_history_entries(session_id, SegmentId::now_v7(), &entries).unwrap(); + assert_eq!(restored, projected); + } + + #[test] + fn annotated_restore_preserves_logical_ids_across_reboot() { + let session_id = SessionId::now_v7(); + let entry = history_entry( + Item::assistant_message("persisted"), + WorkerHistoryProvenance::ModelOutput { + worker: worker_subject(session_id), + }, + ); + let log = vec![LogEntry::AnnotatedSegmentStart { + ts: 0, + session_id, + system_prompt: None, + config: RequestConfig::default(), + history: vec![to_logged_history_entry(&entry)], + forked_from: None, + compacted_from: None, + }]; + let first = restore_history_entries(session_id, SegmentId::now_v7(), &log).unwrap(); + let second = restore_history_entries(session_id, SegmentId::now_v7(), &log).unwrap(); + assert_eq!(first[0].annotation.entry_id, entry.annotation.entry_id); + assert_eq!(second[0].annotation.entry_id, entry.annotation.entry_id); + } + + #[test] + fn compacted_derivation_uses_stable_logical_entry_ids() { + let source = history_entry( + Item::user_message("source"), + WorkerHistoryProvenance::LegacyUnknown, + ); + let summary = HistoryEntry::new( + Item::system_message("summary"), + metadata( + WorkerHistoryProvenance::DerivedSummary, + Some(SessionHistoryDerivation { + sources: vec![source.annotation.entry_id.clone()], + }), + ), + ); + assert_eq!( + summary.annotation.derivation.unwrap().sources, + vec![source.annotation.entry_id] + ); + } +} diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 3d2d8782..21a9371a 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -499,7 +499,10 @@ impl Tool for SubWorkerSpawnTool { InternalWorkerVisibility::ParentClient, Some(child_registry.clone()), Some(Arc::new(move |status| { - if status == InternalWorkerSessionStatus::Failed { + if matches!( + status, + InternalWorkerSessionStatus::Failed | InternalWorkerSessionStatus::Stopped + ) { if let Some(registry) = registry.upgrade() { if let Err(error) = registry.reclaim_internal_scope(&child_name) { tracing::warn!( @@ -1249,7 +1252,7 @@ extract_threshold = 4000 ) .await .unwrap(); - assert!(first_capture.items.iter().any(|item| { + assert!(first_capture.entries.iter().map(|entry| &entry.item).any(|item| { matches!(item, Item::Message { role: Role::Assistant, content, .. } if content.iter().any(|part| matches!(part, ContentPart::Text { text } if text.contains("reviewed")))) })); @@ -1271,7 +1274,7 @@ extract_threshold = 4000 ) .await .unwrap(); - assert!(latest_capture.items.len() > first_capture.items.len()); + assert!(latest_capture.entries.len() > first_capture.entries.len()); fail_requests.store(true, Ordering::SeqCst); send.execute( @@ -1282,16 +1285,16 @@ extract_threshold = 4000 .unwrap(); assert_eq!( record.session.wait_until_idle().await, - InternalWorkerSessionStatus::Failed + InternalWorkerSessionStatus::Stopped ); assert_eq!(calls.load(Ordering::SeqCst), 3); assert!( spawner_scope.snapshot().is_writable(&workspace_root), - "Failed terminal child must release its delegated Workdir session" + "Stopped terminal child must release its delegated Workdir session" ); assert!( !record.workdir_delegation.is_active(), - "failed child must revoke cloned scoped sessions" + "stopped child must revoke cloned scoped sessions" ); assert!(registry.get_internal("reviewer-child").is_some()); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 0b0306e9..28f69623 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -5,16 +5,18 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; -use agen::Item; use agen::llm_client::RequestConfig; use agen::llm_client::client::LlmClient; use agen::llm_client::types::Role; use agen::state::Mutable; -use agen::{Engine, EngineError, EngineResult, ToolOutputLimits, UsageRecord}; +use agen::{ + Engine, EngineError, EngineResult, EngineRunExit, History, HistoryEntry, Item, StopReason, + ToolExecutionPolicy, ToolOutputLimits, UsageRecord, +}; use arc_swap::ArcSwap; use session_store::{ LogEntry, PromptRenderProvenance, SegmentId, SessionExtension, SessionId, Store, StoreError, - SystemItem, segment_log, to_logged, + SystemItem, segment_log, }; use session_store::{ WorkerActiveSegmentRef, WorkerMetadata, WorkerMetadataStore, WorkerReclaimedChild, @@ -23,6 +25,11 @@ use session_store::{ use tracing::{info, warn}; use crate::segment_log_sink::SegmentLogSink; +use crate::session_history::{ + SessionHistoryDerivation, SessionHistoryMetadata, WorkerHistoryProvenance, history_entry, + metadata as new_history_metadata, restore_history_entries, to_logged_history_entry, + worker_subject, +}; use manifest::{ DelegationScope, Permission, ResolveError, Scope, ScopeConfig, ScopeError, ScopeRule, @@ -75,8 +82,8 @@ use crate::skill::{SkillActivationResponse, SkillClientError}; #[cfg(test)] use async_trait::async_trait; use protocol::{ - AlertLevel, AlertSource, CompactionLifecycle, CompactionLifecycleState, Event, RewindSummary, - RewindTarget, RewindTargetId, Segment, + AlertLevel, AlertSource, CompactionLifecycle, CompactionLifecycleState, ErrorCode, Event, + RewindSummary, RewindTarget, RewindTargetId, Segment, }; use tokio::net::UnixStream; use tokio::sync::broadcast; @@ -803,6 +810,70 @@ fn is_ai_materialized_item(item: &Item) -> bool { } } +fn history_annotator( + annotation_writer: LogWriterHandle, + pending_input: Vec>, + pending_committed_history: Arc< + Mutex>>, + >, +) -> impl FnMut(&Item) -> Result +where + St: Store + Clone, +{ + let mut pending_input = std::collections::VecDeque::from(pending_input); + move |item: &Item| { + if let Some(entry) = pending_input.pop_front() { + return Ok(entry.annotation); + } + if let Some(entry) = { + let mut pending = pending_committed_history + .lock() + .expect("pending committed history poisoned"); + pending + .front() + .filter(|entry| entry.item == *item) + .cloned() + .map(|entry| { + pending.pop_front(); + entry + }) + } { + return Ok(entry.annotation); + } + + 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() { + WorkerHistoryProvenance::ModelOutput { worker: subject } + } else { + // Unknown user/system append paths fail closed. Trusted system + // producers must precommit through `SystemItemCommitter`. + WorkerHistoryProvenance::LegacyUnknown + }; + let metadata = new_history_metadata(origin, None); + let entry = session_store::LoggedHistoryEntry { + item: item.clone().into(), + metadata: metadata.clone(), + }; + let log_entry = if item.is_tool_result() { + LogEntry::AnnotatedToolResult { + ts: segment_log::now_millis(), + entry, + } + } else { + LogEntry::AnnotatedAssistantItem { + ts: segment_log::now_millis(), + entry, + } + }; + annotation_writer + .append_entry(log_entry) + .map_err(|error| error.to_string())?; + Ok(metadata) + } +} + /// Cheap-cloneable bundle of (store + shared session pointer + sink) /// handed to the worker callback and the interceptor so they can /// commit `LogEntry` values directly without going through an mpsc @@ -828,8 +899,12 @@ where self.store.append(loc.session_id, loc.segment_id, &entry)?; self.state.increment_entries(); if let Some(in_flight) = &self.in_flight { - if let LogEntry::AssistantItem { item, .. } = &entry { - let item_for_clear = item.clone(); + let committed_item = match &entry { + LogEntry::AssistantItem { item, .. } => Some(item.clone()), + LogEntry::AnnotatedAssistantItem { entry, .. } => Some(entry.item.clone()), + _ => None, + }; + if let Some(item_for_clear) = committed_item { in_flight.clear_for_committed_item_then(&item_for_clear, || { self.sink.publish(entry); }); @@ -856,11 +931,23 @@ where pub trait SystemItemCommitter: Send + Sync { fn commit_log_entry(&self, entry: LogEntry) -> Result<(), StoreError>; - fn commit_system_item(&self, item: SystemItem) -> Result<(), StoreError> { - self.commit_log_entry(LogEntry::SystemItem { + fn commit_system_item( + &self, + item: SystemItem, + ) -> Result, StoreError> { + let metadata = new_history_metadata( + WorkerHistoryProvenance::BackendInstruction { operation_id: None }, + None, + ); + let history_item = item.to_history_item(); + self.commit_log_entry(LogEntry::AnnotatedSystemItem { ts: segment_log::now_millis(), - item, - }) + entry: session_store::LoggedSystemHistoryEntry { + item, + metadata: metadata.clone(), + }, + })?; + Ok(HistoryEntry::new(history_item, metadata)) } } @@ -897,6 +984,64 @@ where pub const WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN: &str = "worker.input-submission.v1"; +#[derive(Clone)] +struct PreparedFlowProjection { + selector: String, + instructions: String, + definition_id: String, + definition_revision: u64, + instance_id: String, + state_id: String, +} + +/// Sole live owner of committed model-visible Worker history. +/// +/// `Engine` borrows this history only while executing a run. The revision is +/// advanced together with every live rewrite so projections can fence stale +/// observations without maintaining a second transcript. +#[derive(Clone)] +pub struct WorkerSession { + session_id: SessionId, + revision: u64, + history: History, +} + +impl WorkerSession { + fn new(session_id: SessionId, entries: Vec>) -> Self { + let revision = u64::try_from(entries.len()).unwrap_or(u64::MAX); + Self { + session_id, + revision, + history: History::from_entries(entries), + } + } + + pub fn session_id(&self) -> SessionId { + self.session_id + } + + pub fn revision(&self) -> u64 { + self.revision + } + + pub fn history(&self) -> &History { + &self.history + } + + fn history_mut(&mut self) -> &mut History { + &mut self.history + } + + fn note_mutation(&mut self) { + self.revision = self.revision.saturating_add(1); + } + + fn replace_history(&mut self, entries: Vec>) { + self.history.replace_entries(entries); + self.note_mutation(); + } +} + /// An independent agent execution unit. /// /// Holds a [`Engine`] directly and persists session state via @@ -904,7 +1049,11 @@ pub const WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN: &str = "worker.input-submiss pub struct Worker { manifest: WorkerManifest, /// Always `Some` outside of `run()`/`resume()`. - engine: Option>, + engine: Option>, + /// Sole live authority for committed model-visible history. + session: WorkerSession, + /// Worker-owned interruption recovery marker. + last_run_interrupted: bool, store: St, /// Optional write-through hook for name-keyed Worker metadata. Production /// constructors install this from the same FsStore that owns the session @@ -1003,6 +1152,10 @@ pub struct Worker { /// drains it and returns `ContinueWith` so the items land in /// history right after the user message that referenced them. pending_attachments: Arc>>, + /// Ephemeral handoff for system items that were durably committed by the + /// interceptor before Agen applies them to live typed history. + pending_committed_history: + Arc>>>, /// Scope allocation in the machine-wide lock file. `Some` for /// Workers built via `from_manifest` / `from_manifest_spawned` / /// `restore_from_manifest` (production paths); `None` for the @@ -1102,11 +1255,14 @@ impl Worker // model is configured. system_prompt / request_config / cache_key // are unused on this path, so we deliberately skip copying them. let source_worker = self.engine.as_ref().expect("worker present"); - let mut worker = Engine::new(source_worker.client().clone()); - worker.set_history(source_worker.history().to_vec()); + let worker = Engine::::new_annotated( + source_worker.client().clone(), + ); Self { manifest: self.manifest.clone(), engine: Some(worker), + session: self.session.clone(), + last_run_interrupted: false, store: self.store.clone(), worker_metadata_writer: None, segment_state: self.segment_state.clone(), @@ -1135,6 +1291,7 @@ impl Worker ai_activity_counter: self.ai_activity_counter.clone(), pending_notifies: NotifyBuffer::new(), pending_attachments: Arc::new(Mutex::new(Vec::::new())), + pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), scope_allocation: None, callback_socket: None, runtime_ticket_role: None, @@ -1153,7 +1310,9 @@ impl Worker log_writer: None, } } +} +impl Worker { /// Build a `LogWriterHandle` carrying everything the worker /// callback / interceptor needs to commit `LogEntry` values /// directly: store handle, the shared session pointer, and the @@ -1199,25 +1358,9 @@ impl Worker /// interrupted-turn prep) before they reach the worker's history, so this /// callback would otherwise double-write them. pub fn wire_history_persistence(&mut self) { - let writer = self.log_writer_handle(); - self.engine_mut().on_history_append(move |item| { - if item.is_user_message() { - return Ok(()); - } - if matches!( - item, - Item::Message { - role: agen::Role::System, - .. - } - ) { - return Ok(()); - } - let entry = session_store::classify_history_item(item, segment_log::now_millis()); - writer - .append_entry(entry) - .map_err(|error| error.to_string()) - }); + // History records are committed by the annotation callback before Agen + // applies the corresponding entry. A second observer callback would + // create an unannotated duplicate and is intentionally not installed. if self.manifest.session.record_event_trace { let writer = self.log_writer_handle(); self.engine_mut() @@ -1253,7 +1396,9 @@ impl Worker } self.history_persistence_wired = true; } +} +impl Worker { pub fn spawn_post_run_memory_jobs(&mut self) { // Drop a finished prior handle so we can spawn a fresh task. // If the prior task is still running, coalesce by skipping — @@ -1275,7 +1420,7 @@ impl Worker } } -impl Worker { +impl Worker { /// Create a new Worker from a pre-built Engine and store. /// /// Callers must pass path-free workspace context separately from explicit @@ -1289,7 +1434,7 @@ impl Worker { /// should parse it themselves and call [`set_system_prompt_template`]. pub async fn new( manifest: WorkerManifest, - worker: Engine, + worker: Engine, store: St, workspace_context: WorkerWorkspaceContext, filesystem_authority: WorkerFilesystemAuthority, @@ -1308,6 +1453,8 @@ impl Worker { let mut worker = Self { manifest, engine: Some(worker), + session: WorkerSession::new(session_id, Vec::new()), + last_run_interrupted: false, store, worker_metadata_writer: None, segment_state: SegmentState::new(session_id, segment_id, 0), @@ -1336,6 +1483,7 @@ impl Worker { ai_activity_counter: Arc::new(AtomicUsize::new(0)), pending_notifies: NotifyBuffer::new(), pending_attachments: Arc::new(Mutex::new(Vec::::new())), + pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), scope_allocation: None, callback_socket: None, runtime_ticket_role: None, @@ -1545,22 +1693,46 @@ impl Worker { /// This deliberately does not scan `.yoi/skills` locally: when a Workspace /// HTTP client is available, catalog/detail/activation authority belongs to /// the Workspace backend API. - pub fn activate_skill(&mut self, name: &str) -> Result { + pub fn activate_skill(&mut self, name: &str) -> Result + where + St: Clone + 'static, + { let activation = self.workspace_client().activate_skill(name)?; self.ensure_segment_head()?; let body = format!( "Agent Skill `{}` activated from {}.\n\n{}", activation.name, activation.provenance.id, activation.body ); - self.commit_entry(LogEntry::SystemItem { + let skill_metadata = new_history_metadata( + WorkerHistoryProvenance::BackendInstruction { operation_id: None }, + None, + ); + self.commit_entry(LogEntry::AnnotatedSystemItem { ts: segment_log::now_millis(), - item: SystemItem::SkillActivation { - name: activation.name.clone(), - body: body.clone(), + entry: session_store::LoggedSystemHistoryEntry { + item: SystemItem::SkillActivation { + name: activation.name.clone(), + body: body.clone(), + }, + metadata: skill_metadata.clone(), }, })?; - self.engine_mut() - .append_history(std::iter::once(agen::Item::system_message(body)))?; + let history_entry = HistoryEntry::new(agen::Item::system_message(body), skill_metadata); + let mut annotate = history_annotator( + self.log_writer_handle(), + vec![history_entry.clone()], + self.pending_committed_history.clone(), + ); + let (engine, session) = ( + self.engine.as_mut().expect("worker present"), + &mut self.session, + ); + engine.append_history_with( + session.history_mut(), + std::iter::once(history_entry.item), + &mut annotate, + )?; + session.note_mutation(); Ok(activation) } @@ -1623,7 +1795,7 @@ impl Worker { } /// Direct access to the underlying Engine. - pub fn engine(&self) -> &Engine { + pub fn engine(&self) -> &Engine { self.engine.as_ref().expect("worker taken during run") } @@ -1631,10 +1803,19 @@ impl Worker { /// /// Use this to register tools, hooks, or subscribers before calling /// [`run`](Self::run). - pub fn engine_mut(&mut self) -> &mut Engine { + pub fn engine_mut(&mut self) -> &mut Engine { self.engine.as_mut().expect("worker taken during run") } + #[cfg(test)] + fn set_history_for_test(&mut self, items: Vec) { + let entries = items + .into_iter() + .map(|item| HistoryEntry::new(item, SessionHistoryMetadata::legacy_unknown())) + .collect(); + self.session.replace_history(entries); + } + /// Install enabled feature modules into the Worker host surfaces. pub fn install_features( &mut self, @@ -1759,13 +1940,15 @@ impl Worker { ))); } - let Some(LogEntry::UserInput { segments, .. }) = entries.get(target.user_input_entry_index) - else { - return Err(RewindError::Invalid( - "rewind target is no longer a user message".into(), - )); + let input = match entries.get(target.user_input_entry_index) { + Some(LogEntry::UserInput { segments, .. }) + | Some(LogEntry::AnnotatedUserInput { segments, .. }) => segments.clone(), + _ => { + return Err(RewindError::Invalid( + "rewind target is no longer a user message".into(), + )); + } }; - let input = segments.clone(); let truncate_entries = rewind_truncate_entries(&entries, target.user_input_entry_index); let retained = entries[..truncate_entries].to_vec(); let tool_side_effect_warning = suffix_has_tool_side_effects(&entries[truncate_entries..]); @@ -1782,13 +1965,17 @@ impl Worker { self.segment_state.set_entries_written(truncate_entries); self.sink.truncate_silent(truncate_entries); - self.task_feature.restore_from_history(&state.history); - let history = state.history; - self.engine_mut().set_history(history); + let history_entries = restore_history_entries(loc.session_id, loc.segment_id, &retained) + .map_err(|error| RewindError::Invalid(error.into()))?; + let projected_history = history_entries + .iter() + .map(|entry| entry.item.clone()) + .collect::>(); + self.task_feature.restore_from_history(&projected_history); + self.session.replace_history(history_entries); self.engine_mut().set_request_config(state.config); self.engine_mut().set_turn_count(state.turn_count); - self.engine_mut() - .set_last_run_interrupted(state.last_run_interrupted); + self.last_run_interrupted = state.last_run_interrupted; self.engine_mut() .set_active_run_turn_count(state.active_run_turn_count); self.user_segments = state.user_segments; @@ -1857,9 +2044,18 @@ impl Worker { self.write_worker_metadata_pending() } - /// Current history items held by the underlying Engine. - pub fn history(&self) -> &[Item] { - self.engine().history() + /// Provider-visible projection of the current typed Worker session history. + /// The authoritative item+provenance entries remain owned by `WorkerSession`. + pub fn history(&self) -> Vec { + self.session.history().items_cloned() + } + + pub fn session_history(&self) -> &History { + self.session.history() + } + + pub fn worker_session(&self) -> &WorkerSession { + &self.session } /// Snapshot of the cumulative LLM Usage measurement timeline. @@ -2163,7 +2359,7 @@ impl Worker { let usage_history_handle = compact_state.as_ref().map(|_| self.usage_history.clone()); - let interceptor = WorkerInterceptor::new( + let interceptor = WorkerInterceptor::new_with_history_queue( registry, compact_state, usage_history_handle, @@ -2171,6 +2367,7 @@ impl Worker { self.pending_attachments.clone(), self.prompts.clone(), self.log_writer.clone(), + self.pending_committed_history.clone(), ) .with_usage_tracker(self.usage_tracker.clone()) .with_prompt_workspace_id( @@ -2293,7 +2490,10 @@ impl Worker { /// Equivalent to `run(vec![Segment::text(s)])`. The dumb-client /// counterpart of [`protocol::Method::run_text`]; primarily for /// tests and tools that have only a string in hand. - pub async fn run_text(&mut self, s: impl Into) -> Result { + pub async fn run_text(&mut self, s: impl Into) -> Result + where + St: Clone + 'static, + { self.run(vec![Segment::text(s)]).await } @@ -2331,7 +2531,10 @@ impl Worker { /// Wires up worker hooks, ensures the session is materialized on the /// store, and runs pre-run compact (joining any in-flight memory task /// first so extract sees a stable history range). - async fn prepare_for_run(&mut self) -> Result<(), WorkerError> { + async fn prepare_for_run(&mut self) -> Result<(), WorkerError> + where + St: Clone + 'static, + { self.refresh_prompt_projection_for_future_operations()?; self.ensure_interceptor_installed(); self.ensure_system_prompt_materialized().await?; @@ -2356,14 +2559,14 @@ impl Worker { .expect("usage_history poisoned") .len(); EmptyTurnRollbackSnapshot { - history_len: self.engine().history().len(), + history_len: self.session.history().len(), user_segments_len: self.user_segments.len(), entries_written: self.segment_state.entries_written(), sink_len: self.sink.len(), pending_attachments, usage_history_len, ai_activity_count: self.ai_activity_counter.load(Ordering::SeqCst), - last_run_interrupted: self.engine().last_run_interrupted(), + last_run_interrupted: self.last_run_interrupted, active_run_turn_count: self.engine().active_run_turn_count(), flow_runtime_state: self .flow_runtime_state @@ -2375,17 +2578,21 @@ impl Worker { fn should_rollback_empty_turn( &self, - result: &Result, + result: &EngineRunExit, snapshot: &EmptyTurnRollbackSnapshot, ) -> bool { - if !matches!(result, Err(EngineError::Cancelled)) { + if !matches!( + result, + EngineRunExit::Paused | EngineRunExit::Interrupted(StopReason::Cancelled) + ) { return false; } if self.ai_activity_counter.load(Ordering::SeqCst) != snapshot.ai_activity_count { return false; } - !self.engine().history()[snapshot.history_len..] + !self.session.history().entries()[snapshot.history_len..] .iter() + .map(|entry| &entry.item) .any(is_ai_materialized_item) } @@ -2393,9 +2600,9 @@ impl Worker { &mut self, snapshot: EmptyTurnRollbackSnapshot, ) -> Result<(), StoreError> { - self.engine_mut().truncate_history(snapshot.history_len); - self.engine_mut() - .set_last_run_interrupted(snapshot.last_run_interrupted); + self.session.history_mut().truncate(snapshot.history_len); + self.session.note_mutation(); + self.last_run_interrupted = snapshot.last_run_interrupted; self.engine_mut() .set_active_run_turn_count(snapshot.active_run_turn_count); *self @@ -2425,7 +2632,14 @@ impl Worker { fn prepare_flow_input( &self, input: Vec, - ) -> Result<(Vec, Option), WorkerError> { + ) -> Result< + ( + Vec, + Option, + Option, + ), + WorkerError, + > { let flow_segments = input .iter() .filter_map(|segment| match segment { @@ -2434,7 +2648,7 @@ impl Worker { }) .collect::>(); if flow_segments.is_empty() { - return Ok((input, None)); + return Ok((input, None, None)); } if flow_segments.len() != 1 { return Err(WorkerError::FlowInput( @@ -2509,16 +2723,15 @@ impl Worker { let (state, initial_instructions) = flow::FlowRuntimeState::start(&source, uuid::Uuid::now_v7().to_string()) .map_err(|error| WorkerError::FlowInput(error.to_string()))?; - let input = input - .into_iter() - .map(|segment| match segment { - Segment::Flow { .. } => Segment::Text { - content: initial_instructions.clone(), - }, - other => other, - }) - .collect(); - Ok((input, Some(state))) + let projection = PreparedFlowProjection { + selector: selector.to_string(), + instructions: initial_instructions, + definition_id: state.instance.definition_id.clone(), + definition_revision: state.instance.definition_revision, + instance_id: state.instance.instance_id.clone(), + state_id: state.instance.current_state.to_string(), + }; + Ok((input, Some(state), Some(projection))) } /// Send user input and run until the LLM turn completes. @@ -2532,16 +2745,40 @@ impl Worker { /// If the between-turns compaction threshold is exceeded mid-run, /// the Engine is aborted, history is compacted, and execution resumes /// automatically. - pub async fn run(&mut self, input: Vec) -> Result { + pub async fn run(&mut self, input: Vec) -> Result + where + St: Clone + 'static, + { self.run_with_input_extensions(input, Vec::new()).await } pub(crate) async fn run_with_input_extensions( + &mut self, + input: Vec, + input_extensions: Vec, + ) -> Result + where + St: Clone + 'static, + { + self.run_with_input_extensions_and_commit_hook(input, input_extensions, || {}) + .await + } + + /// Run user input and invoke `on_input_committed` only after the annotated + /// input has crossed both the durable Store and live SegmentLogSink commit + /// boundaries. The Controller uses this fence before exposing `Running`, so + /// every in-flight snapshot for a user turn includes its committed input. + pub(crate) async fn run_with_input_extensions_and_commit_hook( &mut self, input: Vec, mut input_extensions: Vec, - ) -> Result { - let (input, pending_flow_state) = self.prepare_flow_input(input)?; + on_input_committed: F, + ) -> Result + where + St: Clone + 'static, + F: FnOnce(), + { + let (input, pending_flow_state, flow_projection) = self.prepare_flow_input(input)?; if let Some(state) = pending_flow_state.as_ref() { let payload = serde_json::to_value(state).map_err(|error| { WorkerError::FlowInput(format!("serialize Flow runtime state: {error}")) @@ -2574,13 +2811,18 @@ impl Worker { trigger: protocol::InvokeKind::UserSend, })?; - // Persist the user input as typed segments before the worker - // pushes its flattened copy into history. save_delta deliberately - // skips the resulting `is_user_message()` item to avoid double-write. - self.commit_entry(LogEntry::UserInput { + let projected_input = self.projected_input_history(&input, flow_projection.as_ref()); + + // Persist original typed segments together with the exact ordered + // model-visible item+origin projection before any entry becomes live. + self.commit_entry(LogEntry::AnnotatedUserInput { ts: segment_log::now_millis(), segments: input.clone(), extensions: input_extensions, + history: projected_input + .iter() + .map(to_logged_history_entry) + .collect(), })?; if let Some(state) = pending_flow_state { *self @@ -2589,12 +2831,12 @@ impl Worker { .expect("flow_runtime_state poisoned") = Some(state); } self.user_segments.push(input.clone()); + on_input_committed(); // Resolve `@` file refs to system messages stashed for the // WorkerInterceptor to attach right after the user message. Resolution // failures are non-fatal alerts. let attachments = self.resolve_file_refs(&input).await; - let flattened = self.flatten_segments(&input); if !attachments.is_empty() { *self .pending_attachments @@ -2602,13 +2844,41 @@ impl Worker { .expect("pending_attachments poisoned") = attachments; } - let history_before = self.engine.as_ref().unwrap().history().len(); + let history_before = self.session.history().len(); + let pending_input = projected_input; + let input_entry = pending_input + .last() + .cloned() + .expect("projected Worker input is never empty"); + let prefix_items = pending_input + .iter() + .take(pending_input.len().saturating_sub(1)) + .map(|entry| entry.item.clone()) + .collect::>(); + let input_string = input_entry.item.as_text().unwrap_or_default(); + let mut annotate = history_annotator( + self.log_writer_handle(), + pending_input, + self.pending_committed_history.clone(), + ); + + if !prefix_items.is_empty() { + let (engine, session) = ( + self.engine.as_mut().expect("worker present"), + &mut self.session, + ); + engine + .append_history_with(session.history_mut(), prefix_items, &mut annotate) + .map_err(|error| WorkerError::InvalidState(error.to_string()))?; + } - // lock → run → unlock let worker = self.engine.take().expect("worker taken during run"); - let mut locked = worker.lock(); - let result = locked.run(flattened).await; + let mut locked = worker.lock(self.session.history()); + let result = locked + .run_with_annotation(self.session.history_mut(), input_string, &mut annotate) + .await; self.engine = Some(locked.unlock()); + self.session.note_mutation(); if self.should_rollback_empty_turn(&result, &rollback_snapshot) { self.rollback_empty_turn(rollback_snapshot)?; @@ -2678,50 +2948,80 @@ impl Worker { /// must happen before `prepare_for_run`: proactive compaction checkpoints /// only resumable runs, never the run this invocation is abandoning. fn prepare_interrupted_history_for_fresh_run(&mut self) -> Result<(), WorkerError> { - if self.engine().last_run_interrupted() { + if self.last_run_interrupted { self.apply_interrupt_prep()?; - self.engine_mut().set_last_run_interrupted(false); + self.last_run_interrupted = false; + self.engine_mut().set_active_run_turn_count(None); } Ok(()) } - /// Stage the post-interruption cleanup at the front of worker - /// history: close every unanswered `Item::ToolCall` with a synthetic - /// `Item::ToolResult` (Anthropic wire-validity), then append a - /// system note so the LLM understands the prior turn was cut - /// short. Called from `Worker::run` when the worker's - /// `last_run_interrupted` flag is set (i.e. the Worker just transitioned - /// out of Paused via a new user input). - fn apply_interrupt_prep(&mut self) -> Result<(), WorkerError> { + /// Durably close every unanswered ToolCall before the interrupted run's + /// final lifecycle record/status is published. + fn terminalize_orphan_tool_calls(&mut self) -> Result<(), WorkerError> { let tool_result_summary = self .prompts() .load_full() .interrupt_tool_result_summary() .map_err(WorkerError::from)?; + let history_items = self.history(); + let closures = crate::interrupt_prep::orphan_tool_result_closures( + &history_items, + &tool_result_summary, + ); + if closures.is_empty() { + return Ok(()); + } + + let subject = worker_subject(self.session.session_id()); + for item in closures { + let entry = HistoryEntry::new( + item, + new_history_metadata( + WorkerHistoryProvenance::ToolOutput { + worker: subject.clone(), + }, + None, + ), + ); + self.commit_entry(LogEntry::AnnotatedToolResult { + ts: segment_log::now_millis(), + entry: to_logged_history_entry(&entry), + })?; + self.session.history_mut().push_entry(entry); + self.session.note_mutation(); + } + Ok(()) + } + + fn apply_interrupt_prep(&mut self) -> Result<(), WorkerError> { + self.terminalize_orphan_tool_calls()?; let system_note = self .prompts() .load_full() .interrupt_system_note() .map_err(WorkerError::from)?; - let closures = crate::interrupt_prep::orphan_tool_result_closures( - self.engine().history(), - &tool_result_summary, - ); - if !closures.is_empty() { - self.engine_mut().append_history(closures)?; - } let interrupt_prompt_provenance = self.prompt_render_provenance("internal.interrupt_system_note"); - self.commit_entry(LogEntry::SystemItem { + let interrupt_metadata = new_history_metadata( + WorkerHistoryProvenance::BackendInstruction { operation_id: None }, + None, + ); + self.commit_entry(LogEntry::AnnotatedSystemItem { ts: segment_log::now_millis(), - item: SystemItem::Interrupt { - body: system_note.clone(), - prompt_provenance: Some(interrupt_prompt_provenance), + entry: session_store::LoggedSystemHistoryEntry { + item: SystemItem::Interrupt { + body: system_note.clone(), + prompt_provenance: Some(interrupt_prompt_provenance), + }, + metadata: interrupt_metadata.clone(), }, })?; - self.engine_mut() - .append_history(std::iter::once(agen::Item::system_message(system_note)))?; + let interrupt_entry = + HistoryEntry::new(agen::Item::system_message(system_note), interrupt_metadata); + self.session.history_mut().push_entry(interrupt_entry); + self.session.note_mutation(); Ok(()) } @@ -2733,49 +3033,52 @@ impl Worker { /// The explicit `PausedTurnAbandoned` marker preserves durable lifecycle /// semantics without claiming another `run` / `resume` completed. pub fn cancel_paused_turn(&mut self) -> Result<(), WorkerError> { - if !self.engine().last_run_interrupted() { + if !self.last_run_interrupted { return Ok(()); } self.apply_interrupt_prep()?; - self.engine_mut().set_last_run_interrupted(false); + self.last_run_interrupted = false; self.commit_entry(LogEntry::PausedTurnAbandoned { ts: segment_log::now_millis(), })?; Ok(()) } - /// Flatten a typed segment list into the single string the Engine - /// receives as the user message, and emit user-facing alerts for - /// segments that fall through to placeholder (unknown variants from a newer client). - /// `FileRef` is handled separately by `resolve_file_refs`. The text - /// reconstruction itself comes from `Segment::flatten_to_text`, - /// shared with replay paths that should not re-alert. - fn flatten_segments(&self, segments: &[Segment]) -> String { - for seg in segments { - match seg { - Segment::Text { .. } | Segment::Paste { .. } | Segment::FileRef { .. } => {} - Segment::Flow { selector } => { - self.alert( - AlertLevel::Error, - AlertSource::Worker, - format!( - "received unresolved Flow invocation {selector:?}; Runtime must resolve Flow segments through Workspace authority before Worker input" - ), - ); - } - Segment::Unknown => { - self.alert( - AlertLevel::Warn, - AlertSource::Worker, - "received unknown segment kind from a newer client; \ - passed to LLM as placeholder" - .into(), - ); - } - } + fn projected_input_history( + &self, + input: &[Segment], + flow_projection: Option<&PreparedFlowProjection>, + ) -> Vec> { + if let Some(flow) = flow_projection { + return input + .iter() + .map(|segment| match segment { + Segment::Flow { .. } => history_entry( + Item::user_message(flow.instructions.clone()), + WorkerHistoryProvenance::FlowInstruction { + selector: flow.selector.clone(), + definition_id: flow.definition_id.clone(), + definition_revision: flow.definition_revision, + instance_id: flow.instance_id.clone(), + state_id: flow.state_id.clone(), + }, + ), + other => history_entry( + Item::user_message(Segment::flatten_to_text(std::slice::from_ref(other))), + // Current public submit transport does not carry a + // trusted account/Worker subject envelope. Fail closed + // instead of promoting role=user to HumanInput. + WorkerHistoryProvenance::LegacyUnknown, + ), + }) + .collect(); } - Segment::flatten_to_text(segments) + + vec![history_entry( + Item::user_message(Segment::flatten_to_text(input)), + WorkerHistoryProvenance::LegacyUnknown, + )] } /// Run a turn triggered by `Method::Notify` while the Worker is idle. @@ -2789,7 +3092,10 @@ impl Worker { pub async fn run_for_notification( &mut self, kind: protocol::InvokeKind, - ) -> Result { + ) -> Result + where + St: Clone + 'static, + { debug_assert!( matches!( kind, @@ -2811,27 +3117,43 @@ impl Worker { trigger: kind, })?; - let history_before = self.engine.as_ref().unwrap().history().len(); - + let history_before = self.session.history().len(); + let mut annotate = history_annotator( + self.log_writer_handle(), + Vec::new(), + self.pending_committed_history.clone(), + ); let worker = self.engine.take().expect("worker taken during run"); - let mut locked = worker.lock(); - let result = locked.resume().await; + let mut locked = worker.lock(self.session.history()); + let result = locked + .resume_with_annotation(self.session.history_mut(), &mut annotate) + .await; self.engine = Some(locked.unlock()); + self.session.note_mutation(); self.handle_worker_result(result, history_before).await } /// Resume from a paused state. - pub async fn resume(&mut self) -> Result { + pub async fn resume(&mut self) -> Result + where + St: Clone + 'static, + { self.prepare_for_run().await?; - let history_before = self.engine.as_ref().unwrap().history().len(); - - // lock → resume → unlock + let history_before = self.session.history().len(); + let mut annotate = history_annotator( + self.log_writer_handle(), + Vec::new(), + self.pending_committed_history.clone(), + ); let worker = self.engine.take().expect("worker taken during run"); - let mut locked = worker.lock(); - let result = locked.resume().await; + let mut locked = worker.lock(self.session.history()); + let result = locked + .resume_with_annotation(self.session.history_mut(), &mut annotate) + .await; self.engine = Some(locked.unlock()); + self.session.note_mutation(); self.handle_worker_result(result, history_before).await } @@ -2850,12 +3172,18 @@ impl Worker { let loc = self.segment_state.location(); let entries_written = self.segment_state.entries_written(); if entries_written == 0 { - let initial = LogEntry::SegmentStart { + let initial = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id: loc.session_id, system_prompt: w.get_system_prompt().map(String::from), config: w.request_config().clone(), - history: to_logged(w.history()), + history: self + .session + .history() + .entries() + .iter() + .map(to_logged_history_entry) + .collect(), forked_from: None, compacted_from: None, }; @@ -2880,12 +3208,18 @@ impl Worker { // and is broadcast through the sink so existing subscribers reset // their view. let fork_segment_id = session_store::new_segment_id(); - let entry = LogEntry::SegmentStart { + let entry = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id: loc.session_id, system_prompt: w.get_system_prompt().map(String::from), config: w.request_config().clone(), - history: to_logged(w.history()), + history: self + .session + .history() + .entries() + .iter() + .map(to_logged_history_entry) + .collect(), forked_from: Some(session_store::SegmentOrigin { segment_id: loc.segment_id, at_turn_index: w.turn_count(), @@ -2927,23 +3261,50 @@ impl Worker { /// `Yielded`), so restore remains consistent. async fn handle_worker_result( &mut self, - result: Result, + result: EngineRunExit, history_before: usize, - ) -> Result { + ) -> Result + where + St: Clone + 'static, + { + if matches!(&result, EngineRunExit::Interrupted(_)) { + self.terminalize_orphan_tool_calls()?; + } self.persist_turn(history_before, &result).await?; - if matches!(result, Ok(EngineResult::Yielded)) { + if matches!(result, EngineRunExit::Yielded) { + self.last_run_interrupted = true; return self.do_compact_and_resume().await; } - if result.is_ok() { + if !matches!(result, EngineRunExit::Interrupted(_)) { if let Some(ref state) = self.compact_state { state.set_just_compacted(false); } } - result - .map(WorkerRunResult::from) - .map_err(WorkerError::Engine) + + match result { + EngineRunExit::Finished => { + self.last_run_interrupted = false; + Ok(WorkerRunResult::Finished) + } + EngineRunExit::Paused => { + self.last_run_interrupted = true; + Ok(WorkerRunResult::Paused) + } + EngineRunExit::Interrupted(StopReason::LimitReached) => { + self.last_run_interrupted = false; + Ok(WorkerRunResult::LimitReached) + } + EngineRunExit::Interrupted(reason) => { + self.last_run_interrupted = true; + Ok(WorkerRunResult::Interrupted { + code: stop_reason_error_code(&reason), + message: stop_reason_message(&reason), + }) + } + EngineRunExit::Yielded => unreachable!("yielded handled above"), + } } fn persist_compaction_lifecycle( @@ -2997,7 +3358,10 @@ impl Worker { &mut self, ) -> std::pin::Pin< Box> + Send + '_>, - > { + > + where + St: Clone + 'static, + { Box::pin(async move { // Thrash detection: if we just compacted and hit the threshold again, // something is wrong. @@ -3153,7 +3517,7 @@ impl Worker { async fn persist_turn( &mut self, history_before: usize, - result: &Result, + result: &EngineRunExit, ) -> Result<(), StoreError> { // Per-item commits for AssistantItem / ToolResult / SystemItem // entries are expected to have landed synchronously: the @@ -3169,9 +3533,9 @@ impl Worker { // slice from `history_before` inline so the test's // `restore`-style assertions still see entries on disk. if !self.history_persistence_wired { - let new_items: Vec = self.engine.as_ref().unwrap().history()[history_before..] + let new_items: Vec = self.session.history().entries()[history_before..] .iter() - .cloned() + .map(|entry| entry.item.clone()) .collect(); let ts = segment_log::now_millis(); for item in &new_items { @@ -3251,22 +3615,43 @@ impl Worker { .push(record); } - let interrupted = self.engine.as_ref().unwrap().last_run_interrupted(); + let interrupted = matches!( + result, + EngineRunExit::Paused + | EngineRunExit::Yielded + | EngineRunExit::Interrupted(StopReason::Cancelled) + | EngineRunExit::Interrupted(StopReason::ContextWindowExceeded) + | EngineRunExit::Interrupted(StopReason::Unexpected(_)) + ); let active_run_turn_count = self.engine.as_ref().unwrap().active_run_turn_count(); match result { - Ok(r) => { + EngineRunExit::Finished | EngineRunExit::Paused | EngineRunExit::Yielded => { + let result = match result { + EngineRunExit::Finished => EngineResult::Finished, + EngineRunExit::Paused => EngineResult::Paused, + EngineRunExit::Yielded => EngineResult::Yielded, + EngineRunExit::Interrupted(_) => unreachable!(), + }; self.commit_entry(LogEntry::RunCompleted { ts: segment_log::now_millis(), interrupted, - result: r.clone(), + result, active_run_turn_count, })?; } - Err(e) => { + EngineRunExit::Interrupted(StopReason::LimitReached) => { + self.commit_entry(LogEntry::RunCompleted { + ts: segment_log::now_millis(), + interrupted: false, + result: EngineResult::LimitReached, + active_run_turn_count, + })?; + } + EngineRunExit::Interrupted(reason) => { self.commit_entry(LogEntry::RunErrored { ts: segment_log::now_millis(), interrupted, - message: e.to_string(), + message: stop_reason_message(reason), })?; } } @@ -3355,11 +3740,18 @@ impl Worker { // within `retained_tokens`. Item-granular, turn boundaries ignored. let cut = self.split_for_retained(retained_tokens); - let worker = self.engine.as_ref().expect("worker taken during run"); - let history = worker.history(); - let retain_from = cut.index.min(history.len()); - let retained_items = history[retain_from..].to_vec(); - let items_to_summarise = history[..retain_from].to_vec(); + let history_entries = self.session.history().entries(); + let retain_from = cut.index.min(history_entries.len()); + let retained_history_entries = history_entries[retain_from..].to_vec(); + let retained_items = retained_history_entries + .iter() + .map(|entry| entry.item.clone()) + .collect::>(); + let entries_to_summarise = history_entries[..retain_from].to_vec(); + let items_to_summarise = entries_to_summarise + .iter() + .map(|entry| entry.item.clone()) + .collect::>(); // Compaction-related knobs. Fall through to manifest defaults when // `[compaction]` is omitted entirely. let ( @@ -3487,9 +3879,9 @@ impl Worker { .with_module( crate::feature::builtin::session_explore::SessionExploreFeature::new( crate::feature::builtin::session_explore::SessionExploreState::new( - crate::session_capture::SessionCapture::new( + crate::session_capture::SessionCapture::from_history_entries( self.segment_id().to_string(), - items_to_summarise.clone(), + entries_to_summarise.clone(), ), ), ), @@ -3766,6 +4158,36 @@ impl Worker { max: result_context_max_tokens, }); } + let original_entries = self.session.history().entries(); + let derived_sources = original_entries + .iter() + .map(|entry| entry.annotation.entry_id.clone()) + .collect::>(); + let mut original_cursor = 0usize; + let compacted_history_entries = new_history + .iter() + .cloned() + .map(|item| { + if let Some((offset, original)) = original_entries[original_cursor..] + .iter() + .enumerate() + .find(|(_, original)| original.item == item) + { + original_cursor += offset + 1; + HistoryEntry::new(item, original.annotation.clone()) + } else { + HistoryEntry::new( + item, + new_history_metadata( + WorkerHistoryProvenance::DerivedSummary, + Some(SessionHistoryDerivation { + sources: derived_sources.clone(), + }), + ), + ) + } + }) + .collect::>(); // Build the SegmentStart entry for the new compacted segment. // Inherits the source Segment's session_id so the compacted @@ -3777,12 +4199,15 @@ impl Worker { let old_loc = self.segment_state.location(); let source_turn_count = self.engine.as_ref().unwrap().turn_count(); let w = self.engine.as_ref().unwrap(); - let entry = LogEntry::SegmentStart { + let entry = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id: old_loc.session_id, system_prompt: w.get_system_prompt().map(String::from), config: w.request_config().clone(), - history: to_logged(&new_history), + history: compacted_history_entries + .iter() + .map(to_logged_history_entry) + .collect(), forked_from: None, compacted_from: Some(session_store::SegmentOrigin { segment_id: old_loc.segment_id, @@ -3846,7 +4271,7 @@ impl Worker { self.user_segments.drain(..drop_n); } - self.engine.as_mut().unwrap().set_history(new_history); + self.session.replace_history(compacted_history_entries); // Compaction-introduced system messages are part of the new // SegmentStart's history (broadcast above) — clients derive // their blocks from `SegmentStart.history`. No per-item @@ -4086,12 +4511,7 @@ impl Worker { return Ok(ExtractDecision::Skipped); } - let current_history_len = self - .engine - .as_ref() - .expect("engine present") - .history() - .len(); + let current_history_len = self.session.history().len(); if current_history_len <= processed_history_len { audit .emit( @@ -4184,9 +4604,8 @@ impl Worker { ) .await; - let items_to_extract = self.engine.as_ref().expect("worker present").history() - [processed_history_len..current_history_len] - .to_vec(); + let entries_to_extract = + self.session.history().entries()[processed_history_len..current_history_len].to_vec(); let extract_worker_max_turns = memory_cfg .extract_worker_max_turns @@ -4236,9 +4655,9 @@ impl Worker { segment_id: source_segment_id.to_string(), range: [start_entry as u64, end_entry as u64], }; - let session_view = crate::session_capture::SessionCapture::new( + let session_view = crate::session_capture::SessionCapture::from_history_entries( source_segment_id.to_string(), - items_to_extract, + entries_to_extract, ); let session_explore_state = SessionExploreState::new(session_view.clone()); let memory_extract_state = MemoryExtractState::new( @@ -4467,6 +4886,9 @@ impl Worker { fn extract_internal_worker_lifecycle_error(lifecycle: &WorkerRunResult) -> Option { match lifecycle { WorkerRunResult::RolledBack => Some(WorkerError::Engine(EngineError::Cancelled)), + WorkerRunResult::Interrupted { message, .. } => { + Some(WorkerError::Engine(EngineError::Aborted(message.clone()))) + } WorkerRunResult::Finished | WorkerRunResult::Paused | WorkerRunResult::LimitReached => None, } } @@ -4719,7 +5141,10 @@ where segment_id, )?; - let mut worker = Engine::new(common.client); + let mut worker = + Engine::, Mutable, SessionHistoryMetadata>::new_annotated( + common.client, + ); apply_worker_manifest(&mut worker, &manifest.engine); worker.set_cache_key(Some(segment_id.to_string())); let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store)); @@ -4729,6 +5154,8 @@ where let mut worker = Self { manifest, engine: Some(worker), + session: WorkerSession::new(session_id, Vec::new()), + last_run_interrupted: false, store, worker_metadata_writer, segment_state: SegmentState::new(session_id, segment_id, 0), @@ -4757,6 +5184,7 @@ where ai_activity_counter: Arc::new(AtomicUsize::new(0)), pending_notifies: NotifyBuffer::new(), pending_attachments: Arc::new(Mutex::new(Vec::::new())), + pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), scope_allocation: Some(scope_allocation), callback_socket: None, runtime_ticket_role: None, @@ -4799,7 +5227,10 @@ where } let session_id = session_store::new_session_id(); let segment_id = session_store::new_segment_id(); - let mut engine = Engine::new(common.client); + let mut engine = + Engine::, Mutable, SessionHistoryMetadata>::new_annotated( + common.client, + ); apply_worker_manifest(&mut engine, &manifest.engine); engine.set_cache_key(Some(segment_id.to_string())); let scope = SharedScope::new(common.scope); @@ -4807,6 +5238,8 @@ where let mut worker = Self { manifest, engine: Some(engine), + session: WorkerSession::new(session_id, Vec::new()), + last_run_interrupted: false, store, worker_metadata_writer: None, segment_state: SegmentState::new(session_id, segment_id, 0), @@ -4835,6 +5268,7 @@ where ai_activity_counter: Arc::new(AtomicUsize::new(0)), pending_notifies: NotifyBuffer::new(), pending_attachments: Arc::new(Mutex::new(Vec::::new())), + pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), scope_allocation: None, callback_socket: None, runtime_ticket_role: None, @@ -4910,7 +5344,10 @@ where segment_id, )?; - let mut worker = Engine::new(common.client); + let mut worker = + Engine::, Mutable, SessionHistoryMetadata>::new_annotated( + common.client, + ); apply_worker_manifest(&mut worker, &manifest.engine); worker.set_cache_key(Some(segment_id.to_string())); let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store)); @@ -4920,6 +5357,8 @@ where let mut worker = Self { manifest, engine: Some(worker), + session: WorkerSession::new(session_id, Vec::new()), + last_run_interrupted: false, store, worker_metadata_writer, segment_state: SegmentState::new(session_id, segment_id, 0), @@ -4948,6 +5387,7 @@ where ai_activity_counter: Arc::new(AtomicUsize::new(0)), pending_notifies: NotifyBuffer::new(), pending_attachments: Arc::new(Mutex::new(Vec::::new())), + pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), scope_allocation: Some(scope_allocation), callback_socket: Some(callback_socket), runtime_ticket_role: None, @@ -5198,7 +5638,10 @@ where // Build the worker and apply the manifest defaults first, then // overwrite the pieces the session log is authoritative for. - let mut worker = Engine::new(common.client); + let mut worker = + Engine::, Mutable, SessionHistoryMetadata>::new_annotated( + common.client, + ); apply_worker_manifest(&mut worker, &manifest.engine); worker.set_cache_key(Some(segment_id.to_string())); if let Some(ref prompt) = state.system_prompt { @@ -5208,18 +5651,23 @@ where // (the Worker's one and only write path that prepends a summary at // history[0]). Restoring the anchor lets Anthropic re-use a // stable cache prefix for long-lived restored sessions. + let restored_history_entries = + restore_history_entries(session_id, segment_id, &raw_entries).map_err(|error| { + WorkerError::InvalidState(format!("restore typed Worker session history: {error}")) + })?; + let restored_history = restored_history_entries + .iter() + .map(|entry| entry.item.clone()) + .collect::>(); let anchored_on_summary = matches!( - state.history.first(), + restored_history.first(), Some(Item::Message { role: agen::Role::System, .. }) ); - let restored_history = state.history.clone(); - worker.set_history(restored_history); worker.set_request_config(state.config.clone()); worker.set_turn_count(state.turn_count); - worker.set_last_run_interrupted(state.last_run_interrupted); worker.set_active_run_turn_count(state.active_run_turn_count); if anchored_on_summary { worker.set_cache_anchor(Some(0)); @@ -5234,6 +5682,8 @@ where let mut worker = Self { manifest, engine: Some(worker), + session: WorkerSession::new(session_id, restored_history_entries), + last_run_interrupted: state.last_run_interrupted, store, worker_metadata_writer, segment_state: SegmentState::new(session_id, segment_id, state.entries_count), @@ -5266,6 +5716,7 @@ where ai_activity_counter: Arc::new(AtomicUsize::new(0)), pending_notifies: NotifyBuffer::new(), pending_attachments: Arc::new(Mutex::new(Vec::::new())), + pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), scope_allocation: Some(scope_allocation), callback_socket: None, runtime_ticket_role: None, @@ -5363,9 +5814,21 @@ where /// Note: `system_prompt` is intentionally not applied here. It is a /// minijinja template that is parsed by `Worker::from_manifest` and /// rendered once at first turn in `ensure_system_prompt_materialized`. -pub fn apply_worker_manifest(worker: &mut Engine, wm: &manifest::EngineManifest) { +pub fn apply_worker_manifest( + worker: &mut Engine, + wm: &manifest::EngineManifest, +) { worker.set_request_config(request_config_from_engine_manifest(wm)); worker.set_max_turns(wm.max_turns.map(|n| n.get())); + // 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 + // mechanical per-call terminalization. + worker.set_tool_execution_policy(ToolExecutionPolicy { + pause_safe_boundary_timeout: Duration::from_millis(100), + cancellation_request_timeout: Duration::from_millis(250), + terminal_confirmation_timeout: Duration::from_millis(500), + }); worker.set_tool_output_limits(Some(ToolOutputLimits { default_max_bytes: wm.tool_output.default_max_bytes, per_tool: wm.tool_output.per_tool.clone(), @@ -5470,8 +5933,36 @@ fn restore_manifest_from_worker_metadata_snapshot( } } +fn stop_reason_error_code(reason: &StopReason) -> ErrorCode { + match reason { + StopReason::ContextWindowExceeded | StopReason::Unexpected(EngineError::Client(_)) => { + ErrorCode::ProviderError + } + StopReason::Unexpected(EngineError::Tool(_)) => ErrorCode::ToolError, + StopReason::LimitReached + | StopReason::Cancelled + | StopReason::Unexpected( + EngineError::Aborted(_) + | EngineError::Cancelled + | EngineError::PauseRequested + | EngineError::ConfigWarnings(_) + | EngineError::HistoryAppend(_) + | EngineError::ToolAttemptFence(_), + ) => ErrorCode::Internal, + } +} + +fn stop_reason_message(reason: &StopReason) -> String { + match reason { + StopReason::LimitReached => "engine turn limit reached".to_string(), + StopReason::ContextWindowExceeded => "model context window reached".to_string(), + StopReason::Cancelled => "engine run cancelled".to_string(), + StopReason::Unexpected(error) => format!("unexpected engine failure: {error}"), + } +} + /// Result of a Worker run. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub enum WorkerRunResult { /// The LLM finished its turn normally. Finished, @@ -5479,6 +5970,8 @@ pub enum WorkerRunResult { Paused, /// The worker reached its configured max_turns limit. LimitReached, + /// The run was interrupted by a known or unexpected terminal cause. + Interrupted { code: ErrorCode, message: String }, /// The submit-time user turn was rolled back because no AI output was materialized. RolledBack, } @@ -5763,27 +6256,30 @@ fn build_rewind_targets(segment_id: uuid::Uuid, entries: &[LogEntry]) -> Vec (segments, ts), + _ => continue, + }; + turn_index += 1; + let truncate_entries = rewind_truncate_entries(entries, entry_index); + let tool_warning = suffix_has_tool_side_effects(&entries[truncate_entries..]); + targets.push(RewindTarget { + id: RewindTargetId { + segment_id, + user_input_entry_index: entry_index, + }, + expected_head_entries: head_entries, + truncate_entries, + turn_index, + timestamp_ms: Some(*ts), + preview: preview_segments(segments), + eligible: true, + disabled_reason: None, + warning: tool_warning.then(|| { + "history suffix will be discarded; tool side effects are not undone".into() + }), + }); } targets.reverse(); targets @@ -5804,8 +6300,9 @@ fn rewind_truncate_entries(entries: &[LogEntry], user_input_entry_index: usize) fn suffix_has_tool_side_effects(entries: &[LogEntry]) -> bool { entries.iter().any(|entry| match entry { - LogEntry::ToolResult { .. } => true, + LogEntry::ToolResult { .. } | LogEntry::AnnotatedToolResult { .. } => true, LogEntry::AssistantItem { item, .. } => logged_item_is_tool_call(item), + LogEntry::AnnotatedAssistantItem { entry, .. } => logged_item_is_tool_call(&entry.item), _ => false, }) } @@ -6964,6 +7461,118 @@ mod build_summary_prompt_tests { } } + #[derive(Clone)] + struct PauseResumeClient { + calls: Arc, + } + + impl PauseResumeClient { + fn new() -> Self { + Self { + calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + } + } + } + + #[async_trait::async_trait] + impl agen::llm_client::LlmClient for PauseResumeClient { + async fn stream( + &self, + _request: agen::llm_client::Request, + ) -> Result< + std::pin::Pin< + Box< + dyn futures::Stream< + Item = Result, + > + Send, + >, + >, + agen::llm_client::ClientError, + > { + use agen::llm_client::{Event, ResponseStatus, StatusEvent}; + let call = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let events = if call == 0 { + vec![ + Event::tool_use_start(0, "call_pending", "pending_once"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] + } else { + vec![ + Event::text_block_start(0), + Event::text_delta(0, "done"), + Event::text_block_stop(0, None), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] + }; + Ok(Box::pin(futures::stream::iter(events.into_iter().map(Ok)))) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + } + + #[derive(Clone)] + struct CountingPendingTool { + calls: Arc, + } + + #[async_trait::async_trait] + impl agen::tool::Tool for CountingPendingTool { + async fn execute( + &self, + _input_json: &str, + _ctx: agen::ToolExecutionContext, + ) -> Result { + self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok("executed once".to_string().into()) + } + } + + fn counting_pending_tool( + calls: Arc, + ) -> agen::tool::ToolDefinition { + Arc::new(move || { + let meta = agen::tool::ToolMeta::new("pending_once") + .description("Counts resumable pending execution") + .input_schema(serde_json::json!({"type": "object"})); + ( + meta, + Arc::new(CountingPendingTool { + calls: calls.clone(), + }) as Arc, + ) + }) + } + + #[derive(Clone)] + struct PauseOnceHook { + should_pause: Arc, + } + + #[async_trait::async_trait] + impl crate::hook::Hook for PauseOnceHook { + async fn call( + &self, + _input: &crate::hook::ToolCallSummary, + ) -> crate::hook::HookPreToolAction { + if self + .should_pause + .swap(false, std::sync::atomic::Ordering::SeqCst) + { + crate::hook::HookPreToolAction::Pause + } else { + crate::hook::HookPreToolAction::Continue + } + } + } + #[derive(Clone)] struct NoopClient; @@ -7005,7 +7614,7 @@ mod build_summary_prompt_tests { let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap(); let mut worker = Worker::new( minimal_manifest(), - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store, WorkerWorkspaceContext::no_workspace(), WorkerFilesystemAuthority::None, @@ -7014,12 +7623,12 @@ mod build_summary_prompt_tests { .await .unwrap(); worker.ensure_segment_head().unwrap(); - worker.engine_mut().set_last_run_interrupted(true); + worker.last_run_interrupted = true; worker.engine_mut().set_active_run_turn_count(Some(3)); worker.prepare_interrupted_history_for_fresh_run().unwrap(); - assert!(!worker.engine().last_run_interrupted()); + assert!(!worker.last_run_interrupted); assert_eq!(worker.engine().active_run_turn_count(), None); let checkpoint = active_run_checkpoint_entry( worker.engine().active_run_turn_count(), @@ -7048,7 +7657,7 @@ mod build_summary_prompt_tests { let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap(); let mut worker = Worker::new( minimal_manifest(), - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store, WorkerWorkspaceContext::no_workspace(), WorkerFilesystemAuthority::None, @@ -7058,7 +7667,7 @@ mod build_summary_prompt_tests { .unwrap(); worker.ensure_segment_head().unwrap(); worker.engine_mut().set_turn_count(7); - worker.engine_mut().set_last_run_interrupted(true); + worker.last_run_interrupted = true; worker.engine_mut().set_active_run_turn_count(Some(3)); let session_id = worker.session_id(); @@ -7087,7 +7696,7 @@ mod build_summary_prompt_tests { assert!(matches!( fork_entries.as_slice(), [ - LogEntry::SegmentStart { .. }, + LogEntry::AnnotatedSegmentStart { .. }, LogEntry::ActiveRunCheckpoint { active_turn_count: 3, total_turn_count: 7, @@ -7108,7 +7717,7 @@ mod build_summary_prompt_tests { let workspace_client = Arc::new(RecordingAuditWorkspaceClient::default()); let mut worker = Worker::new( minimal_manifest(), - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store, WorkerWorkspaceContext::with_client( Some(WorkspaceId::new("workspace-test").unwrap()), @@ -7142,7 +7751,7 @@ mod build_summary_prompt_tests { let workspace_client = Arc::new(FlowSourceWorkspaceClient::default()); let mut worker = Worker::new( minimal_manifest(), - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store.clone(), WorkerWorkspaceContext::with_client( Some(WorkspaceId::new("workspace-test").unwrap()), @@ -7177,7 +7786,7 @@ mod build_summary_prompt_tests { }]); assert!(matches!(invalid, Err(WorkerError::FlowInput(_)))); - let (segments, state) = worker + let (segments, state, projection) = worker .prepare_flow_input(vec![ Segment::Flow { selector: "builtin:coder-review".to_string(), @@ -7190,11 +7799,13 @@ mod build_summary_prompt_tests { FLOW_RUNTIME_EXTENSION_DOMAIN, serde_json::to_value(&state).unwrap(), ); + let projected = worker.projected_input_history(&segments, projection.as_ref()); worker - .commit_entry(LogEntry::UserInput { + .commit_entry(LogEntry::AnnotatedUserInput { ts: segment_log::now_millis(), segments: segments.clone(), extensions: vec![extension], + history: projected.iter().map(to_logged_history_entry).collect(), }) .unwrap(); *worker @@ -7202,11 +7813,19 @@ mod build_summary_prompt_tests { .lock() .expect("flow runtime state lock") = Some(state.clone()); - assert_eq!( - segments[0], - Segment::text("Implement the Ticket and request review.") - ); + assert!(matches!( + &segments[0], + Segment::Flow { selector } if selector == "builtin:coder-review" + )); assert_eq!(segments[1], Segment::text("Implement Ticket 00001")); + assert_eq!( + projected[0].item.as_text().as_deref(), + Some("Implement the Ticket and request review.") + ); + assert!(matches!( + projected[0].annotation.origin, + WorkerHistoryProvenance::FlowInstruction { .. } + )); assert_eq!(state.instance.definition_revision, 3); assert_eq!(state.instance.current_state.as_str(), "implement"); assert_eq!(workspace_client.requests.lock().unwrap().len(), 1); @@ -7232,7 +7851,7 @@ mod build_summary_prompt_tests { let mut detached = Worker::new( minimal_manifest(), - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), session_store::FsStore::new(dir.path().join("detached-sessions")).unwrap(), WorkerWorkspaceContext::unavailable( Some(WorkspaceId::new("workspace-test").unwrap()), @@ -7265,7 +7884,7 @@ mod build_summary_prompt_tests { let authority = WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()); let mut worker = Worker::new( manifest, - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store, WorkerWorkspaceContext::local_filesystem(None), authority, @@ -7285,6 +7904,56 @@ mod build_summary_prompt_tests { .unwrap(); } + fn append_annotated_user_turn( + worker: &Worker, + ts: u64, + text: &str, + ) -> Vec { + let user = history_entry( + Item::user_message(text), + WorkerHistoryProvenance::HumanInput { + account_id: "account-1".into(), + }, + ); + let assistant = history_entry( + Item::assistant_message(format!("answer: {text}")), + WorkerHistoryProvenance::ModelOutput { + worker: worker_subject(worker.session.session_id()), + }, + ); + append_test_entry( + worker, + LogEntry::Invoke { + ts, + trigger: protocol::InvokeKind::UserSend, + }, + ); + append_test_entry( + worker, + LogEntry::AnnotatedUserInput { + ts: ts + 1, + segments: vec![Segment::text(text)], + extensions: Vec::new(), + history: vec![to_logged_history_entry(&user)], + }, + ); + append_test_entry( + worker, + LogEntry::AnnotatedAssistantItem { + ts: ts + 2, + entry: to_logged_history_entry(&assistant), + }, + ); + append_test_entry( + worker, + LogEntry::TurnEnd { + ts: ts + 3, + turn_count: 1, + }, + ); + vec![user.annotation, assistant.annotation] + } + fn append_user_turn(worker: &Worker, ts: u64, text: &str) { append_test_entry( worker, @@ -7324,6 +7993,7 @@ mod build_summary_prompt_tests { summary: "wrote a file".into(), content: None, attachments: Vec::new(), + disposition: Default::default(), is_error: false, }, }, @@ -7366,6 +8036,7 @@ mod build_summary_prompt_tests { summary: "wrote a file".into(), content: None, attachments: Vec::new(), + disposition: Default::default(), is_error: false, }, }, @@ -7391,10 +8062,107 @@ mod build_summary_prompt_tests { .len(), expected_truncate_entries ); - assert_eq!(worker.engine().history().len(), 1); + assert_eq!(worker.history().len(), 1); + assert_eq!(worker.history()[0].as_text().unwrap(), "first message"); + } + + #[tokio::test] + async fn annotated_history_rewind_commits_authoritative_prefix() { + let (_dir, mut worker) = rewind_test_worker().await; + let expected_metadata = append_annotated_user_turn(&worker, 10, "first message"); + append_annotated_user_turn(&worker, 20, "second message"); + append_test_entry( + &worker, + LogEntry::AnnotatedToolResult { + ts: 30, + entry: session_store::LoggedHistoryEntry { + item: session_store::LoggedItem::ToolResult { + call_id: "call-v2".into(), + summary: "side effect".into(), + content: None, + attachments: Vec::new(), + disposition: Default::default(), + is_error: false, + }, + metadata: new_history_metadata( + WorkerHistoryProvenance::ToolOutput { + worker: worker_subject(worker.session.session_id()), + }, + None, + ), + }, + }, + ); + + let (head_entries, targets) = worker.list_rewind_targets().unwrap(); + assert_eq!(targets.len(), 2); + assert_eq!(targets[0].preview, "second message"); + assert!(targets[0].truncate_entries > 0); + assert!(targets[0].warning.is_some()); + + let applied = worker + .rewind_to(targets[0].id.clone(), head_entries) + .unwrap(); + assert_eq!(applied.summary.truncated_to_entries, 5); + assert!(matches!( + applied.entries.first(), + Some(LogEntry::AnnotatedSegmentStart { .. }) + )); + let retained_metadata = worker + .session_history() + .entries() + .iter() + .map(|entry| entry.annotation.clone()) + .collect::>(); + assert_eq!(retained_metadata, expected_metadata); + + let source_location = worker.segment_state.location(); + let persisted_prefix = worker + .store + .read_all(source_location.session_id, source_location.segment_id) + .unwrap(); + let restored_prefix = restore_history_entries( + source_location.session_id, + source_location.segment_id, + &persisted_prefix, + ) + .unwrap(); assert_eq!( - worker.engine().history()[0].as_text().unwrap(), - "first message" + restored_prefix + .iter() + .map(|entry| entry.annotation.clone()) + .collect::>(), + expected_metadata + ); + + // Simulate a stale concurrent writer so the next head check forks. + append_test_entry( + &worker, + LogEntry::Extension { + ts: 31, + domain: "test.concurrent-writer".into(), + payload: serde_json::json!({"value": true}), + }, + ); + worker.ensure_segment_head().unwrap(); + let fork_location = worker.segment_state.location(); + assert_ne!(fork_location.segment_id, source_location.segment_id); + let fork_entries = worker + .store + .read_all(fork_location.session_id, fork_location.segment_id) + .unwrap(); + let fork_history = restore_history_entries( + fork_location.session_id, + fork_location.segment_id, + &fork_entries, + ) + .unwrap(); + assert_eq!( + fork_history + .iter() + .map(|entry| entry.annotation.clone()) + .collect::>(), + expected_metadata ); } @@ -7414,7 +8182,61 @@ mod build_summary_prompt_tests { } #[tokio::test] - async fn apply_interrupt_prep_appends_via_callback_and_logs_independent_entries() { + async fn hook_paused_pending_tool_resumes_and_executes_once() { + let dir = tempfile::tempdir().unwrap(); + let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap(); + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let mut engine = + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(PauseResumeClient::new()); + engine.register_tool(counting_pending_tool(calls.clone())); + let mut worker = Worker::new( + minimal_manifest(), + engine, + store, + WorkerWorkspaceContext::no_workspace(), + WorkerFilesystemAuthority::None, + Scope::empty(), + ) + .await + .unwrap(); + let should_pause = Arc::new(std::sync::atomic::AtomicBool::new(true)); + worker.add_pre_tool_call_hook(PauseOnceHook { should_pause }); + + assert_eq!( + worker.run_text("start").await.unwrap(), + WorkerRunResult::Paused + ); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 0); + assert!(worker.history().iter().any(|item| matches!( + item, + Item::ToolCall { call_id, .. } if call_id == "call_pending" + ))); + assert!(!worker.history().iter().any(|item| matches!( + item, + Item::ToolResult { call_id, .. } if call_id == "call_pending" + ))); + + assert_eq!(worker.resume().await.unwrap(), WorkerRunResult::Finished); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!( + worker + .history() + .iter() + .filter(|item| matches!( + item, + Item::ToolResult { + call_id, + disposition: agen::ToolResultDisposition::Success, + .. + } if call_id == "call_pending" + )) + .count(), + 1 + ); + } + + #[tokio::test] + async fn interrupted_result_terminalizes_orphan_before_run_completed() { let dir = tempfile::tempdir().unwrap(); let manifest = minimal_manifest(); let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap(); @@ -7424,7 +8246,7 @@ mod build_summary_prompt_tests { let authority = WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()); let mut worker = Worker::new( manifest, - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store, WorkerWorkspaceContext::local_filesystem(None), authority, @@ -7435,13 +8257,128 @@ mod build_summary_prompt_tests { worker.ensure_segment_head().unwrap(); worker.wire_history_persistence(); - worker - .engine_mut() - .set_history(vec![Item::tool_call("call-1", "Read", "{}")]); + worker.set_history_for_test(vec![ + Item::tool_call("call-known", "Read", "{}"), + Item::tool_result_item_with_disposition_and_attachments( + "call-known", + "known result", + Some("confirmed output".to_string()), + agen::ToolResultDisposition::Success, + Vec::new(), + ), + Item::tool_call("call-orphan", "Bash", "{}"), + ]); + let _ = worker + .handle_worker_result( + EngineRunExit::Interrupted(StopReason::Cancelled), + worker.history().len(), + ) + .await + .unwrap(); + + let history = worker.history(); + assert_eq!( + history + .iter() + .filter(|item| matches!( + item, + Item::ToolResult { + call_id, + disposition: agen::ToolResultDisposition::Success, + .. + } if call_id == "call-known" + )) + .count(), + 1 + ); + assert!(!history.iter().any(|item| matches!( + item, + Item::ToolResult { + call_id, + disposition: agen::ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call-known" + ))); + assert_eq!( + history + .iter() + .filter(|item| matches!( + item, + Item::ToolResult { + call_id, + disposition: agen::ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call-orphan" + )) + .count(), + 1 + ); + + let entries = worker + .store + .read_all( + worker.segment_state.session_id(), + worker.segment_state.segment_id(), + ) + .unwrap(); + let terminal_index = entries + .iter() + .position(|entry| { + matches!( + entry, + LogEntry::AnnotatedToolResult { + entry: session_store::LoggedHistoryEntry { + item: session_store::LoggedItem::ToolResult { + call_id, + disposition: agen::ToolResultDisposition::OutcomeUnknown, + .. + }, + .. + }, + .. + } if call_id == "call-orphan" + ) + }) + .expect("durable OutcomeUnknown closure"); + let final_index = entries + .iter() + .position(|entry| { + matches!( + entry, + LogEntry::RunCompleted { .. } | LogEntry::RunErrored { .. } + ) + }) + .expect("durable final run status"); + assert!(terminal_index < final_index); + } + + #[tokio::test] + async fn apply_interrupt_prep_appends_via_callback_and_logs_independent_entries() { + let dir = tempfile::tempdir().unwrap(); + let manifest = minimal_manifest(); + let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap(); + let cwd = dir.path().join("workspace"); + std::fs::create_dir_all(&cwd).unwrap(); + let scope = Scope::writable(&cwd).unwrap(); + let authority = WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()); + let mut worker = Worker::new( + manifest, + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), + store, + WorkerWorkspaceContext::local_filesystem(None), + authority, + scope, + ) + .await + .unwrap(); + + worker.ensure_segment_head().unwrap(); + worker.wire_history_persistence(); + worker.set_history_for_test(vec![Item::tool_call("call-1", "Read", "{}")]); worker.apply_interrupt_prep().unwrap(); - let history = worker.engine().history(); + let history = worker.history(); assert_eq!(history.len(), 3); assert!(matches!(history[1], Item::ToolResult { ref call_id, .. } if call_id == "call-1")); assert!(matches!( @@ -7465,8 +8402,11 @@ mod build_summary_prompt_tests { .filter(|entry| { matches!( entry, - LogEntry::ToolResult { - item: session_store::LoggedItem::ToolResult { call_id, .. }, + LogEntry::AnnotatedToolResult { + entry: session_store::LoggedHistoryEntry { + item: session_store::LoggedItem::ToolResult { call_id, .. }, + .. + }, .. } if call_id == "call-1" ) @@ -7477,8 +8417,11 @@ mod build_summary_prompt_tests { .filter(|entry| { matches!( entry, - LogEntry::SystemItem { - item: SystemItem::Interrupt { body, .. }, + LogEntry::AnnotatedSystemItem { + entry: session_store::LoggedSystemHistoryEntry { + item: SystemItem::Interrupt { body, .. }, + .. + }, .. } if body == &interrupt_note ) @@ -7500,7 +8443,7 @@ mod build_summary_prompt_tests { let authority = WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()); let mut worker = Worker::new( manifest, - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store, WorkerWorkspaceContext::local_filesystem(None), authority, @@ -7518,15 +8461,15 @@ mod build_summary_prompt_tests { item: dangling_call.clone().into(), }) .unwrap(); - worker.engine_mut().set_history(vec![dangling_call]); - worker.engine_mut().set_last_run_interrupted(true); + worker.set_history_for_test(vec![dangling_call]); + worker.last_run_interrupted = true; worker .run_for_notification(protocol::InvokeKind::Notify) .await .unwrap(); - let history = worker.engine().history(); + let history = worker.history(); assert!(matches!( history.get(1), Some(Item::ToolResult { call_id, .. }) if call_id == "call-1" @@ -7576,7 +8519,7 @@ mod build_summary_prompt_tests { manifest.memory = Some(memory); let mut worker = Worker::new( manifest, - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store, WorkerWorkspaceContext::no_workspace(), WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone()), @@ -7636,7 +8579,7 @@ mod build_summary_prompt_tests { }; let mut worker = Worker::new( manifest, - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store, workspace_context, authority, @@ -7829,7 +8772,7 @@ mod build_summary_prompt_tests { .unwrap() .block_on(Worker::new( manifest, - Engine::new(NoopClient), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient), store, WorkerWorkspaceContext::with_client( Some(WorkspaceId::new("ws-skill").unwrap()), @@ -7867,8 +8810,11 @@ mod build_summary_prompt_tests { assert!(entries.iter().any(|entry| { matches!( entry, - LogEntry::SystemItem { - item: SystemItem::SkillActivation { name, body }, + LogEntry::AnnotatedSystemItem { + entry: session_store::LoggedSystemHistoryEntry { + item: SystemItem::SkillActivation { name, body }, + .. + }, .. } if name == "triage-errors" && body.contains("# Triage Errors") @@ -7899,7 +8845,7 @@ mod build_summary_prompt_tests { let memory_config = manifest.memory.clone().unwrap(); let mut worker = Worker::new( manifest, - Engine::new(client), + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(client), store, WorkerWorkspaceContext::with_client( Some(WorkspaceId::new("workspace-test").unwrap()), @@ -7915,7 +8861,7 @@ mod build_summary_prompt_tests { let evidence = Item::user_message( "The cancellation regression must leave this evidence available for retry.", ); - worker.engine_mut().set_history(vec![evidence.clone()]); + worker.set_history_for_test(vec![evidence.clone()]); worker .commit_entry(LogEntry::UserInput { ts: segment_log::now_millis(), @@ -7974,7 +8920,7 @@ mod build_summary_prompt_tests { .expect("extract pointer lock") .is_none() ); - assert_eq!(worker.engine().history(), &[evidence]); + assert_eq!(worker.history(), &[evidence]); let entries_after = worker .store diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index 3110fd60..67e81367 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -163,7 +163,8 @@ async fn make_worker_with_manifest( let scope = worker::Scope::writable(&pwd).unwrap(); std::mem::forget(pwd_tmp); - let worker = Engine::new(client); + let worker = + Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); let mut worker = Worker::new( manifest, worker, @@ -204,28 +205,34 @@ fn system_texts_in_sink_session_start( ) -> Vec { let (entries, _rx) = worker.sink().subscribe_with_snapshot(); for entry in entries.into_iter().rev() { - if let session_store::LogEntry::SegmentStart { history, .. } = entry { - return history + let history = match entry { + session_store::LogEntry::AnnotatedSegmentStart { history, .. } => history .into_iter() - .filter_map(|logged| { - let item: Item = logged.into(); - match item { - Item::Message { - role: agen::Role::System, - content, - .. - } => Some( - content - .iter() - .map(|p| p.as_text().to_owned()) - .collect::>() - .join(""), - ), - _ => None, - } - }) - .collect(); - } + .map(|entry| entry.item) + .collect::>(), + session_store::LogEntry::SegmentStart { history, .. } => history, + _ => continue, + }; + return history + .into_iter() + .filter_map(|logged| { + let item: Item = logged.into(); + match item { + Item::Message { + role: agen::Role::System, + content, + .. + } => Some( + content + .iter() + .map(|p| p.as_text().to_owned()) + .collect::>() + .join(""), + ), + _ => None, + } + }) + .collect(); } Vec::new() } @@ -337,7 +344,12 @@ permission = "write" // New segment records forked_from pointing at the source. let new_entries = store.read_all(session_id, new_segment_id).unwrap(); match &new_entries[0] { - LogEntry::SegmentStart { + LogEntry::AnnotatedSegmentStart { + session_id: seg_session, + forked_from: Some(origin), + .. + } + | LogEntry::SegmentStart { session_id: seg_session, forked_from: Some(origin), .. diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 3e7e750b..e4ef5b31 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -32,16 +32,29 @@ fn history_from_sink(handle: &WorkerHandle) -> Vec { let mut items = Vec::new(); for entry in entries { match entry { + LogEntry::AnnotatedSegmentStart { history, .. } => { + items.extend(history.into_iter().map(|entry| Item::from(entry.item))); + } LogEntry::SegmentStart { history, .. } => { items.extend(history.into_iter().map(Item::from)); } + LogEntry::AnnotatedUserInput { history, .. } => { + items.extend(history.into_iter().map(|entry| Item::from(entry.item))); + } LogEntry::UserInput { segments, .. } => { let text = protocol::Segment::flatten_to_text(&segments); items.push(Item::user_message(text)); } + LogEntry::AnnotatedAssistantItem { entry, .. } + | LogEntry::AnnotatedToolResult { entry, .. } => { + items.push(Item::from(entry.item)); + } LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => { items.push(Item::from(item)); } + LogEntry::AnnotatedSystemItem { entry, .. } => { + items.push(entry.item.to_history_item()); + } LogEntry::SystemItem { item, .. } => { items.push(item.to_history_item()); } @@ -51,6 +64,14 @@ fn history_from_sink(handle: &WorkerHandle) -> Vec { items } +fn system_item(entry: &LogEntry) -> Option<&session_store::SystemItem> { + match entry { + LogEntry::AnnotatedSystemItem { entry, .. } => Some(&entry.item), + LogEntry::SystemItem { item, .. } => Some(item), + _ => None, + } +} + // --------------------------------------------------------------------------- // Mock LLM Client // --------------------------------------------------------------------------- @@ -192,7 +213,8 @@ async fn make_worker_with_pwd_and_manifest( let scope = manifest::Scope::writable(&pwd).unwrap(); std::mem::forget(pwd_tmp); - let worker = Engine::new(client); + let worker = + Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); let authority = WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone()); let worker = Worker::new( manifest, @@ -784,13 +806,30 @@ async fn snapshot_includes_user_input_for_in_flight_turn() { let client = MockClient::sequential(vec![MockResponse::Hang(simple_text_events())]); let worker = make_worker(client).await; let handle = spawn_controller(worker).await; + let mut events = handle.subscribe(); handle .send(Method::run_text("hello in-flight")) .await .unwrap(); - wait_for_status(&handle, WorkerStatus::Running).await; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + if matches!( + events.recv().await, + Ok(Event::Status { + status: WorkerStatus::Running, + }) + ) { + break; + } + } + }) + .await + .expect("running status event"); + // The Running event is the in-flight visibility fence: the committed + // annotated input must already be available to an immediately attaching + // subscriber rather than racing behind this status transition. let stream = tokio::net::UnixStream::connect(handle.runtime_dir.socket_path()) .await .unwrap(); @@ -804,10 +843,12 @@ async fn snapshot_includes_user_input_for_in_flight_turn() { // Walk the entries, find a `LogEntry::UserInput` and // confirm its segments flatten to our submitted text. let mut found = false; - for value in entries { + for value in &entries { let entry: session_store::LogEntry = - serde_json::from_value(value).expect("LogEntry deserialise"); - if let session_store::LogEntry::UserInput { segments, .. } = entry { + serde_json::from_value(value.clone()).expect("LogEntry deserialise"); + if let session_store::LogEntry::UserInput { segments, .. } + | session_store::LogEntry::AnnotatedUserInput { segments, .. } = entry + { let text = protocol::Segment::flatten_to_text(&segments); if text == "hello in-flight" { found = true; @@ -815,7 +856,10 @@ async fn snapshot_includes_user_input_for_in_flight_turn() { } } } - assert!(found, "snapshot must carry the in-flight UserInput entry"); + assert!( + found, + "snapshot must carry the in-flight UserInput entry: {entries:?}" + ); return; } Event::Alert(_) => continue, @@ -1086,7 +1130,7 @@ async fn run_with_paste_segment_inlines_content_and_emits_typed_user_message() { _ => {} }, entry = entry_rx.recv() => match entry { - Ok(session_store::LogEntry::UserInput { segments, .. }) => { + Ok(session_store::LogEntry::UserInput { segments, .. } | session_store::LogEntry::AnnotatedUserInput { segments, .. }) => { user_input_segments = Some(segments); if saw_turn_end { break; @@ -1317,11 +1361,8 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() { let (entries, _) = handle.sink.subscribe_with_snapshot(); let saw_notify_in_mirror = entries.iter().any(|e| { matches!( - e, - session_store::LogEntry::SystemItem { - item: session_store::SystemItem::Notification { message, .. }, - .. - } if message == "turn finished" + system_item(e), + Some(session_store::SystemItem::Notification { message, .. }) if message == "turn finished" ) }); assert!( @@ -1463,14 +1504,11 @@ async fn worker_event_turn_ended_while_idle_auto_starts_turn_and_injects_system_ let (entries, _) = handle.sink.subscribe_with_snapshot(); let saw_worker_event_in_mirror = entries.iter().any(|e| { matches!( - e, - session_store::LogEntry::SystemItem { - item: session_store::SystemItem::WorkerEvent { - event: protocol::WorkerEvent::TurnEnded { worker_name }, - .. - }, + system_item(e), + Some(session_store::SystemItem::WorkerEvent { + event: protocol::WorkerEvent::TurnEnded { worker_name }, .. - } if worker_name == "child" + }) if worker_name == "child" ) }); assert!( @@ -1552,14 +1590,11 @@ async fn worker_event_scope_sub_delegated_while_idle_stays_control_plane_only() let (entries, _) = handle.sink.subscribe_with_snapshot(); let saw_scope_event_in_mirror = entries.iter().any(|entry| { matches!( - entry, - session_store::LogEntry::SystemItem { - item: session_store::SystemItem::WorkerEvent { - event: protocol::WorkerEvent::ScopeSubDelegated { .. }, - .. - }, + system_item(entry), + Some(session_store::SystemItem::WorkerEvent { + event: protocol::WorkerEvent::ScopeSubDelegated { .. }, .. - } + }) ) }); assert!( @@ -2134,9 +2169,13 @@ async fn paused_then_run_closes_orphan_tool_use_for_next_request() { for item in items { match item { agen::Item::ToolResult { - call_id, summary, .. + call_id, + summary, + disposition, + .. } if call_id == "call_orphan" => { - assert_eq!(summary, "[Interrupted by user]"); + assert_eq!(summary, "Tool execution outcome unknown"); + assert_eq!(*disposition, agen::ToolResultDisposition::OutcomeUnknown); saw_synthetic_tool_result = true; } agen::Item::Message { role, content, .. } if *role == agen::Role::System => { @@ -2327,8 +2366,11 @@ async fn paused_cancel_abandons_resume_and_next_input_is_fresh_run() { assert!( items.iter().any(|item| matches!( item, - agen::Item::ToolResult { call_id, summary, .. } - if call_id == "call_cancelled" && summary == "[Interrupted by user]" + agen::Item::ToolResult { + call_id, + disposition: agen::ToolResultDisposition::OutcomeUnknown, + .. + } if call_id == "call_cancelled" )), "paused cancel should close orphan tool_use before future requests: {items:?}" ); @@ -2373,7 +2415,8 @@ async fn snapshot_contains_user_input(handle: &WorkerHandle, needle: &str) -> bo let entry: session_store::LogEntry = serde_json::from_value(value).expect("LogEntry deserialise"); match entry { - session_store::LogEntry::UserInput { segments, .. } => { + session_store::LogEntry::UserInput { segments, .. } + | session_store::LogEntry::AnnotatedUserInput { segments, .. } => { protocol::Segment::flatten_to_text(&segments).contains(needle) } _ => false, diff --git a/crates/worker/tests/session_metrics_test.rs b/crates/worker/tests/session_metrics_test.rs index 11e46dcc..17e34aa8 100644 --- a/crates/worker/tests/session_metrics_test.rs +++ b/crates/worker/tests/session_metrics_test.rs @@ -188,7 +188,8 @@ async fn make_worker( let pwd = pwd_tmp.path().to_path_buf(); let scope = worker::Scope::writable(&pwd).unwrap(); - let mut worker = Engine::new(client); + let mut worker = + Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); worker.register_tool(big_content_tool_definition(tool_name)); let worker = Worker::new( @@ -460,7 +461,8 @@ async fn metric_write_failure_emits_warn_alert_and_does_not_abort_run() { // protected token budget covers the only user message). That is enough to drive // the failure path: at least one metric attempts to write. let client = MockClient::new(vec![text_response_with_cache("hi", 0, 0)]); - let worker = Engine::new(client); + let worker = + Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); let mut worker = Worker::new( manifest, worker, @@ -536,7 +538,8 @@ permission = "write" let pwd_tmp = tempfile::tempdir().unwrap(); let pwd = pwd_tmp.path().to_path_buf(); let scope = worker::Scope::writable(&pwd).unwrap(); - let worker = Engine::new(client); + let worker = + Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); let mut worker = Worker::new( manifest, worker, diff --git a/crates/worker/tests/system_prompt_template_test.rs b/crates/worker/tests/system_prompt_template_test.rs index d417db11..53ca4141 100644 --- a/crates/worker/tests/system_prompt_template_test.rs +++ b/crates/worker/tests/system_prompt_template_test.rs @@ -130,7 +130,8 @@ async fn make_worker_with_body( EffectivePromptCatalog::new(templates, 1, "test-schema", "test-toolchain").unwrap(); let loader = PromptCatalogSource::builtins_only().with_effective_catalog(projection); - let worker = Engine::new(client); + let worker = + Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); let mut worker = Worker::new( manifest, worker, diff --git a/crates/workspace-server/src/authority.rs b/crates/workspace-server/src/authority.rs index 1d01aad7..dd87052c 100644 --- a/crates/workspace-server/src/authority.rs +++ b/crates/workspace-server/src/authority.rs @@ -22,7 +22,7 @@ use crate::records::{ TicketEvidenceEvent, TicketEvidenceSummary, TicketListPageRequest, TicketMergeRequestSummary, TicketQueryItem, TicketQueryRequest, TicketQueryResponse, TicketRelationView, TicketRoleAssignmentSummary, TicketShowRequest, TicketSummary, TicketSummaryPage, - summarize_body, truncate_body, validate_project_id, + summarize_body, truncate_body, }; use crate::store::{ ControlPlaneStore, MemoryDocumentRecord, MemoryStagingRecord, MemoryStagingResolutionRecord, @@ -633,7 +633,15 @@ impl SqliteWorkspaceAuthority { predicates.push(format!("o.updated_at<{value}")); } if let Some(value) = &query.linked_ticket_id { - let value = bind(SqlValue::Text(value.clone())); + let resolved = self + .store + .resolve_resource_reference( + &self.workspace_id, + WorkspaceResourceKind::Ticket, + value, + )? + .ok_or_else(|| invalid_objective_error("linked Ticket was not found"))?; + let value = bind(SqlValue::Text(resolved)); predicates.push(format!("EXISTS (SELECT 1 FROM objective_ticket_links link WHERE link.workspace_id=o.workspace_id AND link.objective_id=o.objective_id AND link.ticket_id={value})")); } let relevance_rank = if let Some(text) = @@ -1237,6 +1245,10 @@ impl ObjectiveAuthority for SqliteWorkspaceAuthority { .into_iter() .map(|link| link.ticket_id) .collect::>(); + let linked_ticket_keys = linked_tickets + .iter() + .map(|ticket_id| self.resource_key(WorkspaceResourceKind::Ticket, ticket_id)) + .collect::>>()?; let body_md = record.body_md.clone(); let objective = ObjectiveSummary { resource_key: self @@ -1253,6 +1265,7 @@ impl ObjectiveAuthority for SqliteWorkspaceAuthority { items.push(objective_query_item( objective, linked_tickets, + linked_ticket_keys, query.query.as_deref(), &body_md, )); @@ -1339,9 +1352,19 @@ impl ObjectiveAuthority for SqliteWorkspaceAuthority { fn create_objective(&self, input: ObjectiveCreateInput) -> Result { validate_objective_title(&input.title)?; validate_objective_state(&input.state)?; - for ticket_id in &input.linked_tickets { - validate_project_id(ticket_id)?; - } + let linked_tickets = input + .linked_tickets + .iter() + .map(|ticket_reference| { + self.store + .resolve_resource_reference( + &self.workspace_id, + WorkspaceResourceKind::Ticket, + ticket_reference, + )? + .ok_or_else(|| invalid_objective_error("linked Ticket was not found")) + }) + .collect::>>()?; let now = now_rfc3339(); let objective_id = allocate_record_id( unix_epoch_millis_now().map_err(|err| { @@ -1367,8 +1390,7 @@ impl ObjectiveAuthority for SqliteWorkspaceAuthority { updated_at: now.clone(), }; self.store.upsert_objective(&record)?; - let links = input - .linked_tickets + let links = linked_tickets .into_iter() .map(|ticket_id| ObjectiveTicketLinkRecord { workspace_id: self.workspace_id.clone(), @@ -2283,12 +2305,18 @@ fn ticket_query_item( .iter() .map(|objective| objective.id.clone()) .collect(), + linked_objective_keys: detail + .linked_objectives + .iter() + .map(|objective| objective.resource_key.clone()) + .collect(), relation_count: detail.relations.outgoing.len() + detail.relations.incoming.len(), blocker_count: detail.relations.blockers.len(), unresolved_blocker_count: detail.relations.blockers.len(), unresolved_review_count: usize::from(detail.evidence.unresolved_request_changes), evidence: detail.evidence.clone(), merge_request: detail.merge_request.clone(), + current_coder: detail.current_coder.clone(), } } @@ -2399,6 +2427,7 @@ fn ticket_item_after_cursor( fn objective_query_item( objective: ObjectiveSummary, linked_tickets: Vec, + linked_ticket_keys: Vec, text: Option<&str>, body_md: &str, ) -> ObjectiveQueryItem { @@ -2426,6 +2455,7 @@ fn objective_query_item( snippet, linked_ticket_count: linked_tickets.len(), linked_tickets, + linked_ticket_keys, } } @@ -3308,16 +3338,21 @@ VALUES ('workspace-test', 'ticket', 4); assert!(!objective.revision.is_empty()); assert_eq!(objective.linked_ticket_summaries[0].id, "00000000001J2"); assert_eq!(objective.linked_ticket_summaries[0].state, "ready"); + let linked_ticket_key = objective.linked_ticket_summaries[0].resource_key.clone(); let objective_query = authority .query_objectives(ObjectiveQueryRequest { query: Some("Control plane".to_string()), - linked_ticket_id: Some("00000000001J2".to_string()), + linked_ticket_id: Some(linked_ticket_key.clone()), limit: Some(1), ..ObjectiveQueryRequest::default() }) .unwrap(); assert_eq!(objective_query.items.len(), 1); assert_eq!(objective_query.items[0].linked_ticket_count, 1); + assert_eq!( + objective_query.items[0].linked_ticket_keys, + vec![linked_ticket_key] + ); assert_eq!(objective_query.page.limit, 1); let body_query = authority .query_objectives(ObjectiveQueryRequest { @@ -3410,7 +3445,7 @@ VALUES ('workspace-test', 'ticket', 3); title: "Create Objective".to_string(), body_md: "Alpha body".to_string(), state: "active".to_string(), - linked_tickets: vec!["00000000001J2".to_string()], + linked_tickets: vec!["T-1".to_string()], }) .unwrap(); assert_eq!(created.title, "Create Objective"); @@ -3436,14 +3471,14 @@ VALUES ('workspace-test', 'ticket', 3); assert_eq!(state.state, "paused"); assert_eq!( authority - .link_objective_ticket(&created.id, "00000000001J3") + .link_objective_ticket(&created.id, "T-2") .unwrap() .linked_tickets, vec!["00000000001J2", "00000000001J3"] ); assert_eq!( authority - .unlink_objective_ticket(&created.id, "00000000001J2") + .unlink_objective_ticket(&created.id, "T-1") .unwrap() .linked_tickets, vec!["00000000001J3"] diff --git a/crates/workspace-server/src/records.rs b/crates/workspace-server/src/records.rs index 91b85100..c5312fa3 100644 --- a/crates/workspace-server/src/records.rs +++ b/crates/workspace-server/src/records.rs @@ -1,12 +1,9 @@ -use project_record::validate_record_id; use serde::{Deserialize, Serialize}; pub use workspace_api::{ ObjectiveDetail, ObjectiveEventDetail, ObjectiveLinkedTicketSummary, ObjectiveResourceSummary, ObjectiveSummary, QueryPage, }; -use crate::{Error, Result}; - const SUMMARY_BODY_LIMIT: usize = 240; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -387,12 +384,16 @@ pub struct TicketQueryItem { pub snippet: Option, pub matching_event: Option, pub linked_objective_ids: Vec, + #[ts(skip)] + pub linked_objective_keys: Vec, pub relation_count: usize, pub blocker_count: usize, pub unresolved_blocker_count: usize, pub unresolved_review_count: usize, pub evidence: TicketEvidenceSummary, pub merge_request: Option, + #[ts(skip)] + pub current_coder: Option, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -435,6 +436,7 @@ pub struct ObjectiveQueryItem { pub snippet: Option, pub linked_ticket_count: usize, pub linked_tickets: Vec, + pub linked_ticket_keys: Vec, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -521,10 +523,6 @@ mod typescript_tests { } } -pub(crate) fn validate_project_id(id: &str) -> Result<()> { - validate_record_id(id).map_err(|_| Error::InvalidRecordId(id.to_string())) -} - pub(crate) fn summarize_body(body: &str) -> String { let summary = body .lines() diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index e46b9d00..cc374820 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -327,8 +327,6 @@ fn repository_local_path(source: &workspace_api::RepositorySource) -> Option = std::sync::LazyLock::new(|| { @@ -7257,25 +7255,21 @@ fn dispatch_orchestrator_queue_attention(api: &WorkspaceApi) { return; } - let shown = queued - .iter() - .take(ORCHESTRATOR_ATTENTION_TICKET_LIMIT) - .map(|ticket| { - format!( - "- {} — {}", - bounded_orchestrator_attention_text(&ticket.id, 80), - bounded_orchestrator_attention_text(&ticket.title, 240) - ) - }) - .collect::>() - .join("\n"); - let omitted = queued - .len() - .saturating_sub(ORCHESTRATOR_ATTENTION_TICKET_LIMIT); - let omitted_line = if omitted == 0 { - String::new() - } else { - format!("Additional queued Tickets omitted from this notice: {omitted}\n") + let attention_context = match orchestrator_queue_attention_context( + &api.config.workspace_id, + &api.config.workspace_id, + &queued, + ) { + Ok(context) => context, + Err(error) => { + tracing::warn!( + workspace_id = %api.config.workspace_id, + candidate_count = queued.len(), + diagnostic = error, + "orchestrator backlog attention projection rejected" + ); + return; + } }; let Ok(Some(config_state)) = api .config_store @@ -7292,16 +7286,19 @@ fn dispatch_orchestrator_queue_attention(api: &WorkspaceApi) { let Ok(catalog) = worker::PromptCatalog::from_projection(projection.catalog().clone()) else { return; }; - let content = match catalog.render_serializable( - ORCHESTRATOR_ATTENTION_PROMPT_NAME, - &BTreeMap::from([ - ("omitted_line", omitted_line.as_str()), - ("workspace_id", api.config.workspace_id.as_str()), - ("ticket_lines", shown.as_str()), - ]), + let content = match catalog.orchestrator_queue_attention( + worker::OrchestratorQueueAttentionPrompt::Server, + &attention_context, ) { Ok(content) => content, - Err(_) => return, + Err(error) => { + tracing::warn!( + workspace_id = %api.config.workspace_id, + diagnostic = %error, + "orchestrator backlog attention rendering failed" + ); + return; + } }; let accepted = api .runtime @@ -7321,20 +7318,26 @@ fn dispatch_orchestrator_queue_attention(api: &WorkspaceApi) { } } -fn bounded_orchestrator_attention_text(input: &str, max_chars: usize) -> String { - let mut output = String::new(); - for (index, character) in input.chars().enumerate() { - if index == max_chars { - output.push('…'); - break; - } - output.push(if character.is_control() { - ' ' - } else { - character - }); +fn orchestrator_queue_attention_context( + expected_workspace_id: &str, + candidate_workspace_id: &str, + tickets: &[ticket::TicketSummary], +) -> std::result::Result { + if candidate_workspace_id != expected_workspace_id { + return Err("foreign_workspace_ticket_projection"); } - output + let tickets = tickets + .iter() + .map(|ticket| { + let resource_key = ticket + .resource_key + .clone() + .ok_or("missing_ticket_resource_key")?; + worker::OrchestratorQueueAttentionTicket::new(resource_key, ticket.title.clone()) + .map_err(|_| "invalid_ticket_resource_key") + }) + .collect::, _>>()?; + Ok(worker::OrchestratorQueueAttentionContext::new(tickets)) } fn require_online_workspace_orchestrator_source( @@ -9258,9 +9261,8 @@ fn cleanup_working_directory_for_runtime( result.diagnostics, )); }; - let record = workdir_record_from_summary(&api, runtime_id, &working_directory.summary); - api.store.upsert_workdir_registry(&record)?; let mut summary = working_directory.summary; + persist_workdir_cleanup_observation(&api, runtime_id, &summary)?; apply_workdir_occupancy_projection(&api, &mut summary)?; Ok(Json(BrowserWorkingDirectoryDetailResponse { workspace_id: api.config.workspace_id.clone(), @@ -14151,11 +14153,7 @@ fn sync_runtime_workdir_observations( api.store.upsert_workdir_registry(&updated)?; } } else { - record.materialization_status = - workdir_status_from_runtime_miss(result.diagnostics.as_slice()).to_string(); - record.cleanliness = "unknown".to_string(); - record.updated_at = now_registry_timestamp(); - api.store.upsert_workdir_registry(&record)?; + persist_workdir_runtime_miss(api, record, result.diagnostics.as_slice())?; } } Err(_) => { @@ -14169,17 +14167,54 @@ fn sync_runtime_workdir_observations( Ok(response.diagnostics) } -fn workdir_status_from_runtime_miss(diagnostics: &[RuntimeDiagnostic]) -> &'static str { - if diagnostics +fn persist_workdir_cleanup_observation( + api: &WorkspaceApi, + runtime_id: &str, + summary: &WorkingDirectorySummary, +) -> ApiResult<()> { + if summary.status == WorkingDirectoryStatusKind::NotFound { + api.store.delete_workdir_registry( + &api.config.workspace_id, + summary.working_directory_id.as_str(), + )?; + } else { + let record = workdir_record_from_summary(api, runtime_id, summary); + api.store.upsert_workdir_registry(&record)?; + } + Ok(()) +} + +fn workdir_runtime_miss_is_not_found(diagnostics: &[RuntimeDiagnostic]) -> bool { + diagnostics .iter() .any(|diagnostic| diagnostic.code == "working_directory_not_found") - { +} + +fn workdir_status_from_runtime_miss(diagnostics: &[RuntimeDiagnostic]) -> &'static str { + if workdir_runtime_miss_is_not_found(diagnostics) { "not_found" } else { "unknown" } } +fn persist_workdir_runtime_miss( + api: &WorkspaceApi, + mut record: WorkdirRegistryRecord, + diagnostics: &[RuntimeDiagnostic], +) -> ApiResult<()> { + if workdir_runtime_miss_is_not_found(diagnostics) { + api.store + .delete_workdir_registry(&api.config.workspace_id, record.workdir_id.as_str())?; + } else { + record.materialization_status = "unknown".to_string(); + record.cleanliness = "unknown".to_string(); + record.updated_at = now_registry_timestamp(); + api.store.upsert_workdir_registry(&record)?; + } + Ok(()) +} + fn sync_all_runtime_workdir_observations(api: &WorkspaceApi) -> Vec { let mut diagnostics = Vec::new(); let runtimes = api.runtime.list_runtimes(api.config.max_records.min(200)); @@ -16977,22 +17012,24 @@ mod tests { #[test] fn workdir_runtime_miss_uses_exact_typed_code() { + let typed_not_found = [RuntimeDiagnostic { + code: "working_directory_not_found".to_string(), + severity: DiagnosticSeverity::Warning, + message: "missing".to_string(), + }]; assert_eq!( - workdir_status_from_runtime_miss(&[RuntimeDiagnostic { - code: "working_directory_not_found".to_string(), - severity: DiagnosticSeverity::Warning, - message: "missing".to_string(), - }]), + workdir_status_from_runtime_miss(&typed_not_found), "not_found" ); - assert_eq!( - workdir_status_from_runtime_miss(&[RuntimeDiagnostic { - code: "some_other_not_found".to_string(), - severity: DiagnosticSeverity::Warning, - message: "not a typed workdir miss".to_string(), - }]), - "unknown" - ); + assert!(workdir_runtime_miss_is_not_found(&typed_not_found)); + + let unrelated = [RuntimeDiagnostic { + code: "some_other_not_found".to_string(), + severity: DiagnosticSeverity::Warning, + message: "not a typed workdir miss".to_string(), + }]; + assert_eq!(workdir_status_from_runtime_miss(&unrelated), "unknown"); + assert!(!workdir_runtime_miss_is_not_found(&unrelated)); } struct DeterministicExecutionBackend { @@ -19615,7 +19652,8 @@ mod tests { #[tokio::test] async fn orchestrator_running_to_idle_recovers_queued_ticket_without_notification_memory() { let dir = tempfile::tempdir().unwrap(); - let api = test_api(dir.path()).await; + init_clean_git_workspace(dir.path()); + let (api, execution) = test_api_with_recording_backend(dir.path()).await; let backend = browser_ticket_backend(&api).unwrap(); let mut input = ticket::NewTicket::new("Recover queued work"); input.workflow_state = Some(TicketWorkflowState::Queued); @@ -19634,6 +19672,8 @@ mod tests { .await .unwrap(); assert!(started.online); + let startup_inputs = execution.take_inputs(); + assert_eq!(startup_inputs.len(), 1); assert_eq!( api.orchestrator_attention_fingerprint .lock() @@ -19669,6 +19709,76 @@ mod tests { .as_deref(), Some(ticket_ref.id.as_str()) ); + let notifications = execution.take_inputs(); + assert_eq!(notifications.len(), 1); + assert_eq!(notifications[0].0.worker_id.to_string(), worker_id); + let content = ¬ifications[0].1; + assert!(content.starts_with("Queued Tickets require attention:")); + assert!( + content.contains(&format!( + "- {} — Recover queued work", + ticket_ref.resource_key.as_deref().unwrap() + )), + "unexpected notification body: {content:?}" + ); + assert!(content.contains("Reread the current Ticket state before acting")); + assert!(!content.contains(ticket_ref.id.as_str())); + assert!(!content.contains(TEST_WORKSPACE_ID)); + assert!(!content.contains("bounded")); + assert!(!content.contains("omitted")); + + let candidates = backend + .list(ticket::TicketListQuery::states([ + ticket::TicketListState::Queued, + ])) + .unwrap(); + let mut truncated_candidates = (1..=worker::OrchestratorQueueAttentionContext::MAX_TICKETS + + 1) + .map(|index| { + let mut candidate = candidates[0].clone(); + candidate.id = format!("opaque-{index}"); + candidate.resource_key = Some(format!("T-{index}")); + candidate.title = format!("Queued {index}"); + candidate + }) + .collect::>(); + let truncated = orchestrator_queue_attention_context( + TEST_WORKSPACE_ID, + TEST_WORKSPACE_ID, + &truncated_candidates, + ) + .unwrap(); + let rendered = worker::PromptCatalog::builtins_only() + .unwrap() + .orchestrator_queue_attention( + worker::OrchestratorQueueAttentionPrompt::Server, + &truncated, + ) + .unwrap(); + assert!(rendered.contains("- T-20 — Queued 20")); + assert!(!rendered.contains("T-21")); + assert!(rendered.contains("were omitted from this notice: 1")); + assert!(!rendered.contains("opaque-")); + + truncated_candidates[0].resource_key = None; + assert_eq!( + orchestrator_queue_attention_context( + TEST_WORKSPACE_ID, + TEST_WORKSPACE_ID, + &truncated_candidates + ) + .unwrap_err(), + "missing_ticket_resource_key" + ); + assert_eq!( + orchestrator_queue_attention_context( + TEST_WORKSPACE_ID, + "foreign-workspace", + &truncated_candidates + ) + .unwrap_err(), + "foreign_workspace_ticket_projection" + ); } #[tokio::test] @@ -21232,6 +21342,87 @@ mod tests { .unwrap(); } + #[tokio::test] + async fn confirmed_runtime_miss_removes_registry_record_but_unknown_is_retained() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + seed_cleanup_workdir(&api, "deleted-workdir", "present", "clean"); + let deleted = api + .store + .get_workdir_registry(TEST_WORKSPACE_ID, "deleted-workdir") + .unwrap() + .unwrap(); + + persist_workdir_runtime_miss( + &api, + deleted, + &[RuntimeDiagnostic { + code: "working_directory_not_found".to_string(), + severity: DiagnosticSeverity::Warning, + message: "missing".to_string(), + }], + ) + .unwrap(); + + assert!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, "deleted-workdir") + .unwrap() + .is_none() + ); + + seed_cleanup_workdir(&api, "unknown-workdir", "present", "clean"); + let unknown = api + .store + .get_workdir_registry(TEST_WORKSPACE_ID, "unknown-workdir") + .unwrap() + .unwrap(); + persist_workdir_runtime_miss( + &api, + unknown, + &[RuntimeDiagnostic { + code: "runtime_unavailable".to_string(), + severity: DiagnosticSeverity::Warning, + message: "temporarily unavailable".to_string(), + }], + ) + .unwrap(); + assert_eq!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, "unknown-workdir") + .unwrap() + .unwrap() + .materialization_status, + "unknown" + ); + } + + #[tokio::test] + async fn cleanup_not_found_observation_removes_registry_record() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + let working_directory_id = "cleanup-existing"; + seed_cleanup_workdir(&api, working_directory_id, "present", "clean"); + let record = api + .store + .get_workdir_registry(TEST_WORKSPACE_ID, working_directory_id) + .unwrap() + .unwrap(); + let mut summary = workdir_summary_from_record(&record); + summary.status = WorkingDirectoryStatusKind::NotFound; + + persist_workdir_cleanup_observation(&api, "runtime-test", &summary).unwrap(); + + assert!( + api.store + .get_workdir_registry(TEST_WORKSPACE_ID, working_directory_id) + .unwrap() + .is_none() + ); + } + fn seed_cleanup_link(api: &WorkspaceApi, runtime_worker_id: &str, workdir_id: &str) { let runtime_worker_id = runtime_worker_id.parse::().unwrap(); api.store diff --git a/resources/prompts/internal/workspace_orchestrator_queue_attention.md b/resources/prompts/internal/workspace_orchestrator_queue_attention.md index 3d96603b..eee42a88 100644 --- a/resources/prompts/internal/workspace_orchestrator_queue_attention.md +++ b/resources/prompts/internal/workspace_orchestrator_queue_attention.md @@ -1,7 +1,8 @@ -Workspace Orchestrator attention: authoritative Ticket state still contains queued work after the previous turn or after Server recovery. - -Workspace: {{workspace_id}} -Remaining queued Tickets (bounded): -{{ticket_lines}} -{{omitted_line}} -Reread the listed Tickets, their relations, orchestration plans, current assignments, Workers, and Workdirs before acting. Continue only work already authorized by the human `ready -> queued` transition. Do not drain the queue automatically and do not create duplicate assignments, Workers, Workdirs, or merges. If no Ticket is currently actionable, record the durable waiting reason on the authoritative Ticket or orchestration plan and stop. For an actionable queued Ticket, call the guarded `SpawnTicketCoder` operation without first changing Ticket state; it records `queued -> inprogress` only after the Coder, initial input, current assignment, and Workdir finalization are durably accepted. +Queued Tickets require attention: +{% for ticket in tickets -%} +- {{ ticket.resource_key }} {{ separator }} {{ ticket.title }} +{% endfor -%} +{% if omitted_ticket_count > 0 -%} +Additional queued Tickets were omitted from this notice: {{ omitted_ticket_count }}. Re-query current Ticket authority for the complete set. +{% endif -%} +Reread the current Ticket state before acting. Preserve the human queue gate and current assignment, dependency, Worker, and Workdir authority; do not create duplicate work. diff --git a/resources/prompts/panel/orchestrator_idle_queue_notice.md b/resources/prompts/panel/orchestrator_idle_queue_notice.md index 2cbf02fd..eee42a88 100644 --- a/resources/prompts/panel/orchestrator_idle_queue_notice.md +++ b/resources/prompts/panel/orchestrator_idle_queue_notice.md @@ -1,22 +1,8 @@ -Workspace Dashboard observed that this Orchestrator Worker is idle while queued Ticket work is present. - -This is bounded attention only, not scheduler authority. Do not drain the queue automatically. Verify the Ticket is still `queued`, then use the guarded `SpawnTicketCoder` operation without a separate state transition; that operation records `queued -> inprogress` only after Worker creation, initial input, assignment, and Workdir finalization are durably accepted. - -Workspace: {{ workspace }} - -Actionable queued Tickets: -{% for ticket in actionable_tickets -%} -- {{ ticket.id }} — {{ ticket.title }} [{{ ticket.classification }}] +Queued Tickets require attention: +{% for ticket in tickets -%} +- {{ ticket.resource_key }} {{ separator }} {{ ticket.title }} {% endfor -%} - -{% if waiting_tickets | length > 0 -%} -Queued Tickets retained in the session work set but currently waiting: -{% for ticket in waiting_tickets -%} -- {{ ticket.id }} — {{ ticket.title }} [{{ ticket.classification }}]: {{ ticket.waiting_reason }} -{% endfor -%} -{% endif -%} {% if omitted_ticket_count > 0 -%} -Additional queued Tickets omitted from this bounded notice: {{ omitted_ticket_count }} +Additional queued Tickets were omitted from this notice: {{ omitted_ticket_count }}. Re-query current Ticket authority for the complete set. {% endif -%} - -Preserve the existing human gate, dependency/conflict/capacity/dirty-workspace checks, and duplicate-start checks using actual Ticket state, role/session claims, visible Workers, and worktrees. +Reread the current Ticket state before acting. Preserve the human queue gate and current assignment, dependency, Worker, and Workdir authority; do not create duplicate work. diff --git a/web/workspace/deno.lock b/web/workspace/deno.lock index a68fd449..25ae28a6 100644 --- a/web/workspace/deno.lock +++ b/web/workspace/deno.lock @@ -9,9 +9,9 @@ "npm:@codemirror/view@6.43.8": "6.43.8", "npm:@fontsource/ibm-plex-mono@5.3.0": "5.3.0", "npm:@lezer/highlight@1.2.3": "1.2.3", - "npm:@sveltejs/adapter-static@3.0.9": "3.0.9_@sveltejs+kit@2.49.4__@sveltejs+vite-plugin-svelte@6.2.1___svelte@5.45.6___vite@7.2.7__svelte@5.45.6__typescript@5.9.3__vite@7.2.7_@sveltejs+vite-plugin-svelte@6.2.1__svelte@5.45.6__vite@7.2.7_svelte@5.45.6_typescript@5.9.3_vite@7.2.7", - "npm:@sveltejs/kit@2.49.4": "2.49.4_@sveltejs+vite-plugin-svelte@6.2.1__svelte@5.45.6__vite@7.2.7_svelte@5.45.6_typescript@5.9.3_vite@7.2.7", - "npm:@sveltejs/vite-plugin-svelte@6.2.1": "6.2.1_svelte@5.45.6_vite@7.2.7", + "npm:@sveltejs/adapter-static@3.0.9": "3.0.9_@sveltejs+kit@2.49.4__@sveltejs+vite-plugin-svelte@6.2.1___svelte@5.45.6___vite@7.2.7____yaml@2.9.0___yaml@2.9.0__svelte@5.45.6__typescript@5.9.3__vite@7.2.7___yaml@2.9.0__yaml@2.9.0_yaml@2.9.0", + "npm:@sveltejs/kit@2.49.4": "2.49.4_@sveltejs+vite-plugin-svelte@6.2.1__svelte@5.45.6__vite@7.2.7___yaml@2.9.0__yaml@2.9.0_svelte@5.45.6_typescript@5.9.3_vite@7.2.7__yaml@2.9.0_yaml@2.9.0", + "npm:@sveltejs/vite-plugin-svelte@6.2.1": "6.2.1_svelte@5.45.6_vite@7.2.7__yaml@2.9.0_yaml@2.9.0", "npm:clsx@2.1.1": "2.1.1", "npm:cookie@0.6.0": "0.6.0", "npm:decodal-codemirror@0.3.0": "0.3.0_@codemirror+language@6.12.4_@codemirror+view@6.43.8_@lezer+highlight@1.2.3_@lezer+lr@1.4.10", @@ -23,7 +23,8 @@ "npm:svelte-check@4.3.4": "4.3.4_svelte@5.45.6_typescript@5.9.3", "npm:svelte@5.45.6": "5.45.6", "npm:typescript@5.9.3": "5.9.3", - "npm:vite@7.2.7": "7.2.7" + "npm:vite@7.2.7": "7.2.7_yaml@2.9.0", + "npm:yaml@2.9.0": "2.9.0" }, "jsr": { "@std/assert@1.0.19": { @@ -433,13 +434,13 @@ "acorn" ] }, - "@sveltejs/adapter-static@3.0.9_@sveltejs+kit@2.49.4__@sveltejs+vite-plugin-svelte@6.2.1___svelte@5.45.6___vite@7.2.7__svelte@5.45.6__typescript@5.9.3__vite@7.2.7_@sveltejs+vite-plugin-svelte@6.2.1__svelte@5.45.6__vite@7.2.7_svelte@5.45.6_typescript@5.9.3_vite@7.2.7": { + "@sveltejs/adapter-static@3.0.9_@sveltejs+kit@2.49.4__@sveltejs+vite-plugin-svelte@6.2.1___svelte@5.45.6___vite@7.2.7____yaml@2.9.0___yaml@2.9.0__svelte@5.45.6__typescript@5.9.3__vite@7.2.7___yaml@2.9.0__yaml@2.9.0_yaml@2.9.0": { "integrity": "sha512-aytHXcMi7lb9ljsWUzXYQ0p5X1z9oWud2olu/EpmH7aCu4m84h7QLvb5Wp+CFirKcwoNnYvYWhyP/L8Vh1ztdw==", "dependencies": [ "@sveltejs/kit" ] }, - "@sveltejs/kit@2.49.4_@sveltejs+vite-plugin-svelte@6.2.1__svelte@5.45.6__vite@7.2.7_svelte@5.45.6_typescript@5.9.3_vite@7.2.7": { + "@sveltejs/kit@2.49.4_@sveltejs+vite-plugin-svelte@6.2.1__svelte@5.45.6__vite@7.2.7___yaml@2.9.0__yaml@2.9.0_svelte@5.45.6_typescript@5.9.3_vite@7.2.7__yaml@2.9.0_yaml@2.9.0": { "integrity": "sha512-JFtOqDoU0DI/+QSG8qnq5bKcehVb3tCHhOG4amsSYth5/KgO4EkJvi42xSAiyKmXAAULW1/Zdb6lkgGEgSxdZg==", "dependencies": [ "@standard-schema/spec", @@ -465,7 +466,7 @@ ], "bin": true }, - "@sveltejs/vite-plugin-svelte-inspector@5.0.2_@sveltejs+vite-plugin-svelte@6.2.1__svelte@5.45.6__vite@7.2.7_svelte@5.45.6_vite@7.2.7": { + "@sveltejs/vite-plugin-svelte-inspector@5.0.2_@sveltejs+vite-plugin-svelte@6.2.1__svelte@5.45.6__vite@7.2.7___yaml@2.9.0__yaml@2.9.0_svelte@5.45.6_vite@7.2.7__yaml@2.9.0_yaml@2.9.0": { "integrity": "sha512-TZzRTcEtZffICSAoZGkPSl6Etsj2torOVrx6Uw0KpXxrec9Gg6jFWQ60Q3+LmNGfZSxHRCZL7vXVZIWmuV50Ig==", "dependencies": [ "@sveltejs/vite-plugin-svelte", @@ -474,7 +475,7 @@ "vite" ] }, - "@sveltejs/vite-plugin-svelte@6.2.1_svelte@5.45.6_vite@7.2.7": { + "@sveltejs/vite-plugin-svelte@6.2.1_svelte@5.45.6_vite@7.2.7__yaml@2.9.0_yaml@2.9.0": { "integrity": "sha512-YZs/OSKOQAQCnJvM/P+F1URotNnYNeU3P2s4oIpzm1uFaqUEqRxUB0g5ejMjEb5Gjb9/PiBI5Ktrq4rUUF8UVQ==", "dependencies": [ "@sveltejs/vite-plugin-svelte-inspector", @@ -966,7 +967,7 @@ "vfile-message" ] }, - "vite@7.2.7": { + "vite@7.2.7_yaml@2.9.0": { "integrity": "sha512-ITcnkFeR3+fI8P1wMgItjGrR10170d8auB4EpMLPqmx6uxElH3a/hHGQabSHKdqd4FXWO1nFIp9rRn7JQ34ACQ==", "dependencies": [ "esbuild", @@ -974,14 +975,18 @@ "picomatch", "postcss", "rollup", - "tinyglobby" + "tinyglobby", + "yaml" ], "optionalDependencies": [ "fsevents" ], + "optionalPeers": [ + "yaml" + ], "bin": true }, - "vitefu@1.1.2_vite@7.2.7": { + "vitefu@1.1.2_vite@7.2.7__yaml@2.9.0_yaml@2.9.0": { "integrity": "sha512-zpKATdUbzbsycPFBN71nS2uzBUQiVnFoOrr2rvqv34S1lcAgMKKkjWleLGeiJlZ8lwCXvtWaRn7R3ZC16SYRuw==", "dependencies": [ "vite" @@ -993,6 +998,10 @@ "w3c-keyname@2.2.8": { "integrity": "sha512-dpojBhNsCNN7T82Tm7k26A6G9ML3NkhDsnw9n/eoxSRlVBB4CEtIQ/KTCLI2Fwf3ataSXRhYFkQi3SlnFwPvPQ==" }, + "yaml@2.9.0": { + "integrity": "sha512-2AvhNX3mb8zd6Zy7INTtSpl1F15HW6Wnqj0srWlkKLcpYl/gMIMJiyuGq2KeI2YFxUPjdlB+3Lc10seMLtL4cA==", + "bin": true + }, "zimmerframe@1.1.4": { "integrity": "sha512-B58NGBEoc8Y9MWWCQGl/gq9xBCe4IiKM0a2x7GZdQKOW5Exr8S1W24J6OgM1njK8xCRGvAJIL/MxXHf6SkmQKQ==" }, @@ -1024,7 +1033,8 @@ "packageJson": { "dependencies": [ "npm:@fontsource/ibm-plex-mono@5.3.0", - "npm:gen-interface-jp@0.8.0" + "npm:gen-interface-jp@0.8.0", + "npm:yaml@2.9.0" ] } } diff --git a/web/workspace/package.json b/web/workspace/package.json index 63ae3ddf..21d7a9c3 100644 --- a/web/workspace/package.json +++ b/web/workspace/package.json @@ -5,6 +5,7 @@ "type": "module", "dependencies": { "@fontsource/ibm-plex-mono": "5.3.0", - "gen-interface-jp": "0.8.0" + "gen-interface-jp": "0.8.0", + "yaml": "2.9.0" } } diff --git a/web/workspace/src/lib/generated/protocol.ts b/web/workspace/src/lib/generated/protocol.ts index b7383f87..6a357acb 100644 --- a/web/workspace/src/lib/generated/protocol.ts +++ b/web/workspace/src/lib/generated/protocol.ts @@ -8,7 +8,7 @@ export type AlertSource = "worker" | "engine" | "compactor" | "agents_md"; export type CompletionKind = "file"; -export type WorkerStatus = "idle" | "running" | "paused"; +export type WorkerStatus = "idle" | "running" | "paused" | "stopped"; export type TurnResult = "finished" | "paused"; @@ -16,6 +16,8 @@ export type InvokeKind = "user_send" | "notify" | "worker_event" | "system_remin export type RunResult = "finished" | "paused" | "limit_reached" | "rolled_back"; +export type ToolResultDisposition = "success" | "error" | "interrupted" | "cancelled" | "outcome_unknown"; + export type ErrorCode = "already_running" | "not_running" | "not_paused" | "provider_error" | "tool_error" | "invalid_request" | "internal"; export type Permission = "read" | "write"; @@ -191,7 +193,7 @@ summary: string, * Full tool output. Absent when the tool chose to return * summary-only, or when the result was pruned. */ -output?: string | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { entries: Array, greeting: Greeting, status: WorkerStatus, +output?: string | null, disposition?: ToolResultDisposition | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { entries: Array, greeting: Greeting, status: WorkerStatus, /** * Unfinished model output that has already streamed in the current * run but is not yet represented by committed snapshot entries. diff --git a/web/workspace/src/lib/workspace/console/ConsoleLineItem.svelte b/web/workspace/src/lib/workspace/console/ConsoleLineItem.svelte index c2e3a640..a53e3fa3 100644 --- a/web/workspace/src/lib/workspace/console/ConsoleLineItem.svelte +++ b/web/workspace/src/lib/workspace/console/ConsoleLineItem.svelte @@ -8,6 +8,7 @@ }; let { item }: Props = $props(); + let detailOpen = $state(false); let nowMs = $state(Date.now()); $effect(() => { @@ -55,22 +56,20 @@ line.kind !== 'activity' && line.kind !== 'task_reminder' && line.kind !== 'run_stats'; } - function toolSummary(line: ConsoleLine): { label: string; suffix: string; rest: string } { - const [firstLine = '', ...rest] = line.body.split('\n'); - const [label, suffix = ''] = firstLine.split(' — ', 2); - return { - label, - suffix, - rest: rest.join('\n') - }; + function toolLabel(line: ConsoleLine): string { + return line.toolCallLabel ?? line.toolCall?.name ?? line.title; + } + + function toolStatus(line: ConsoleLine): string { + return line.toolStatus ?? line.toolCall?.state ?? ''; } function shouldRenderMarkdown(line: ConsoleLine): boolean { return line.kind === 'user' || line.kind === 'assistant' || line.kind === 'system'; } - function bodyTextAfterToolSummary(line: ConsoleLine): string { - return toolSummary(line).rest; + function toolBodyText(line: ConsoleLine): string { + return detailOpen ? (line.expandedBody ?? line.body) : line.body; } @@ -107,20 +106,27 @@ {:else if item.kind === 'tool'}
- {toolSummary(item).label} - - {toolSummary(item).suffix} + {toolLabel(item)} + {toolStatus(item)} + {#if item.detail} + + {/if}
{/if} {#if item.compaction} {:else if item.kind === 'tool'} - {#if bodyTextAfterToolSummary(item)} + {#if toolBodyText(item)}

{#if isBashTool(item)} - + {:else} - {bodyTextAfterToolSummary(item)} + {toolBodyText(item)} {/if}

{/if} @@ -147,11 +153,10 @@ {/each} {/if} - {#if item.detail} -
- detail + {#if item.detail && detailOpen} +

{item.detail}

-
+ {/if} @@ -306,48 +311,46 @@ .tool-summary { display: flex; align-items: baseline; - gap: 0; + gap: 0.5rem; color: var(--text-muted); font-size: 0.88rem; font-weight: 750; } .tool-label { - flex: 0 0 auto; - color: var(--tui-cyan); - white-space: nowrap; - } - - .tool-separator { - flex: 0 0 auto; - white-space: nowrap; - } - - .tool-suffix { flex: 1 1 auto; min-width: 0; - overflow-wrap: anywhere; + overflow: hidden; + color: var(--tui-cyan); + text-overflow: ellipsis; + white-space: nowrap; } - .tool-separator, - .tool-suffix { + .tool-status { + flex: 0 0 auto; color: var(--tui-dark-gray); + font-size: 0.72rem; + white-space: nowrap; } - .tool-state-error .tool-suffix { + .tool-state-error .tool-status { color: var(--tui-red); } - .tool-state-running .tool-suffix, - .tool-state-streaming_args .tool-suffix, - .tool-state-pending .tool-suffix { + .tool-state-running .tool-status, + .tool-state-streaming_args .tool-status, + .tool-state-pending .tool-status { color: var(--tui-yellow); } - .tool-state-done .tool-suffix { + .tool-state-done .tool-status { color: var(--tui-dark-gray); } + .console-line.error-line .tool-status { + color: var(--tui-red); + } + .message-heading { display: flex; align-items: center; @@ -412,13 +415,47 @@ color: var(--code); } + .tool-detail-button { + margin-inline-start: auto; + border: 1px solid var(--line); + border-radius: 0.35rem; + padding: 0.08rem 0.35rem; + background: var(--bg-raised); + color: var(--text-muted); + cursor: pointer; + font: inherit; + font-size: 0.68rem; + font-weight: 750; + opacity: 0; + pointer-events: none; + transition: opacity 120ms ease; + } + + .console-line:hover .tool-detail-button, + .tool-detail-button:focus-visible, + .tool-detail-button[aria-expanded='true'] { + opacity: 1; + pointer-events: auto; + } + .message-detail { + margin-top: 0.35rem; + border-left: 2px solid var(--line); + padding-left: 0.6rem; color: var(--text-muted); font-size: 0.84rem; } - .message-detail summary { - cursor: pointer; - font-weight: 800; + .message-detail p { + margin: 0; + overflow-wrap: anywhere; + white-space: pre-wrap; + } + + @media (hover: none) { + .tool-detail-button { + opacity: 1; + pointer-events: auto; + } } diff --git a/web/workspace/src/lib/workspace/console/model.test.ts b/web/workspace/src/lib/workspace/console/model.test.ts index 98fb085a..39732416 100644 --- a/web/workspace/src/lib/workspace/console/model.test.ts +++ b/web/workspace/src/lib/workspace/console/model.test.ts @@ -313,26 +313,25 @@ Deno.test("projectConsole groups tool call lifecycle into one Call block", () => !toolLines[0].streaming, "completed tool call should not remain streaming", ); - assert( - toolLines[0].body.includes("$ pwd"), - "Bash command should be summarized", - ); + assertEquals(toolLines[0].toolCallLabel, "Bash($ pwd)"); + assertEquals(toolLines[0].toolStatus, "done"); assert( toolLines[0].body.includes("/repo"), "tool result should be folded into the Call block", ); assert( toolLines[0].body.includes("line9"), - "Bash result preview should include the ninth output line", + "Bash preview should include the ninth output line", ); assert( !toolLines[0].body.includes("line10") && - !toolLines[0].body.includes("line12"), - "Bash result preview should be capped at ten display lines", + toolLines[0].body.includes("… +3 more lines"), + "Bash preview should retain its line cap", ); assert( - toolLines[0].body.includes("… +3 more lines"), - "Bash result preview should show omitted output count", + toolLines[0].expandedBody?.includes("line12") === true && + !toolLines[0].expandedBody?.includes("more lines"), + "Bash detail should show every returned output line", ); assert( toolLines[0].detail?.includes("id: call-1"), @@ -421,7 +420,8 @@ Deno.test("projectConsole streams distinct Bash stdout and stderr through termin ]); const [line] = projection.lines.filter((line) => line.kind === "tool"); - assert(line.body.includes("Bash — failed (exit 7)"), line.body); + assertEquals(line.toolCallLabel, "Bash($ long-command)"); + assertEquals(line.toolStatus, "failed (exit 7)"); assert(!line.body.includes("elapsed"), line.body); assert(!line.body.includes("stdout:"), line.body); assert(line.body.includes("ready\n"), line.body); @@ -463,7 +463,8 @@ Deno.test("snapshot restores bounded in-flight Bash command output", () => { const projection = projectConsole([{ eventId: "snapshot-command", event: snapshot }]); const [line] = projection.lines.filter((line) => line.kind === "tool"); - assert(line.body.includes("Bash — running…"), line.body); + assertEquals(line.toolCallLabel, "Bash($ slow)"); + assertEquals(line.toolStatus, "running…"); assert(!line.body.includes("elapsed"), line.body); assert(!line.body.includes("stdout:"), line.body); assert(line.body.includes("[… earlier stdout omitted]\ntail\n"), line.body); @@ -474,7 +475,7 @@ Deno.test("snapshot restores bounded in-flight Bash command output", () => { assertEquals(line.streaming, true); }); -Deno.test("projectConsole caps default tool request and result previews", () => { +Deno.test("projectConsole caps default preview but keeps complete detail body", () => { const projection = projectConsole([ { eventId: "70", @@ -508,19 +509,100 @@ Deno.test("projectConsole caps default tool request and result previews", () => const [line] = projection.lines.filter((line) => line.kind === "tool"); assertEquals(line.title, "Call · CustomTool"); - assertEquals(line.body.split("\n").length, 7); - assert(line.body.includes("CustomTool — done"), "tool state should be shown"); + assertEquals(line.toolCallLabel, 'CustomTool("first":"one","second":"two","third":"three","fourth":"four")'); + assertEquals(line.toolStatus, "done"); + assertEquals(line.body.split("\n").length, 3); assert( - line.body.includes('"first": "one"'), - "request preview should be shown", + line.body.includes("out1") && line.body.includes("… +3 more lines"), + "normal display should retain the capped response preview", + ); + assert(!line.body.includes("first"), "request arguments should stay in the Call signature and detail"); + assert( + line.detail?.includes("arguments:\nfirst: one") === true && + line.detail?.includes("fourth: four") === true, + "detail metadata should render complete request arguments as YAML", + ); + assert( + line.expandedBody?.includes("out5") === true && + !line.expandedBody?.includes("more lines"), + "detail body should contain the complete result", ); - assert(line.body.includes("out1"), "result preview should be shown"); - assert(!line.body.includes("third"), "request preview should be capped"); - assert(!line.body.includes("out3"), "result preview should be capped"); - assert(line.body.includes("… +"), "overflow marker should be shown"); }); -Deno.test("projectConsole shows Grep query and caps result preview to five entries", () => { +Deno.test("projectConsole renders JSON tool responses as YAML", () => { + const projection = projectConsole([ + { + eventId: "json-call", + event: { + event: "tool_call_done", + data: { + id: "json-tool", + name: "CustomTool", + arguments: "{}", + }, + } satisfies Event, + }, + { + eventId: "json-result", + event: { + event: "tool_result", + data: { + id: "json-tool", + summary: "json completed", + output: JSON.stringify({ + status: "ok", + items: [{ id: 1 }, { id: 2 }], + }), + is_error: false, + }, + } satisfies Event, + }, + { + eventId: "invalid-json-call", + event: { + event: "tool_call_done", + data: { + id: "invalid-json-tool", + name: "CustomTool", + arguments: "{}", + }, + } satisfies Event, + }, + { + eventId: "invalid-json-result", + event: { + event: "tool_result", + data: { + id: "invalid-json-tool", + summary: "invalid json", + output: '{"status": broken}', + is_error: false, + }, + } satisfies Event, + }, + ]); + + const toolLines = projection.lines.filter((line) => line.kind === "tool"); + const jsonLine = toolLines.find((line) => line.id.includes("json-tool")); + const invalidLine = toolLines.find((line) => line.id.includes("invalid-json-tool")); + assert(jsonLine, "JSON tool line should be projected"); + assert(invalidLine, "invalid JSON tool line should be projected"); + assert( + jsonLine.expandedBody?.includes("status: ok") === true && + jsonLine.expandedBody?.includes(" - id: 2") === true, + "detail body should serialize parsed JSON as YAML", + ); + assert( + jsonLine.body.includes("more lines"), + "normal preview should cap the pretty-printed JSON", + ); + assert( + invalidLine.expandedBody?.includes('{"status": broken}') === true, + "invalid JSON-looking output should remain unchanged", + ); +}); + +Deno.test("projectConsole caps Grep preview but keeps complete detail body", () => { const projection = projectConsole([ { eventId: "72", @@ -549,17 +631,19 @@ Deno.test("projectConsole shows Grep query and caps result preview to five entri const [line] = projection.lines.filter((line) => line.kind === "tool"); assertEquals(line.title, "Call · Grep"); - assert( - line.body.includes("Grep — 6 matches"), - "Grep summary should be shown", - ); - assert(line.body.includes("query: needle"), "Grep query should be shown"); + assertEquals(line.toolCallLabel, "Grep(needle)"); + assertEquals(line.toolStatus, "done"); assert(line.body.includes("hit1"), "first result should be shown"); assert(line.body.includes("hit5"), "fifth result should be shown"); - assert(!line.body.includes("hit6"), "sixth result should be capped"); + assert(!line.body.includes("hit6"), "normal preview should retain its result cap"); assert( line.body.includes("… +1 more results"), - "overflow marker should be shown", + "preview should show the omitted result count", + ); + assert( + line.expandedBody?.includes("hit6") === true && + !line.expandedBody?.includes("more results"), + "detail body should show every Grep result", ); }); @@ -594,17 +678,15 @@ Deno.test("projectConsole keeps Grep error detail in the body", () => { const [line] = projection.lines.filter((line) => line.kind === "tool"); assertEquals(line.title, "Call · Grep"); - assert( - line.body.includes("Grep — Failed"), - "error suffix should stay short", - ); + assertEquals(line.toolCallLabel, "Grep(needle)"); + assertEquals(line.toolStatus, "error"); assert( line.body.includes(message), "error detail should remain visible in the body", ); assert( - !line.body.includes(`Grep — ${message}`), - "error detail should not be repeated in the suffix", + !line.toolCallLabel?.includes(message), + "error detail should not be repeated in the Call signature", ); }); @@ -879,9 +961,10 @@ Deno.test("projectConsole keeps streaming tool call updates in the same Call blo assertEquals(toolLines.length, 1); assertEquals(toolLines[0].title, "Call · Read"); assert(toolLines[0].streaming, "streaming tool call should remain streaming"); + assertEquals(toolLines[0].toolCallLabel, "Read(1 file)"); + assertEquals(toolLines[0].toolStatus, "reading…"); assert( - toolLines[0].body.includes("/tmp/a.md") && - toolLines[0].body.includes("Read — reading"), + toolLines[0].body.includes("/tmp/a.md"), "Read call should render aggregate progress and path without content", ); }); @@ -1018,10 +1101,8 @@ Deno.test("projectConsole aggregates Read calls without showing file content", ( const toolLines = projection.lines.filter((line) => line.kind === "tool"); assertEquals(toolLines.length, 1); assertEquals(toolLines[0].title, "Call · Read"); - assert( - toolLines[0].body.includes("Read — 2 files read"), - "aggregate count should be shown", - ); + assertEquals(toolLines[0].toolCallLabel, "Read(2 files)"); + assertEquals(toolLines[0].toolStatus, "done"); assert( toolLines[0].body.includes("/tmp/a.md"), "first path should be listed", @@ -1073,7 +1154,9 @@ Deno.test("projectConsole renders Edit calls with structured diff lines", () => const [line] = projection.lines.filter((line) => line.kind === "tool"); assertEquals(line.title, "Call · Edit"); - assert(line.body.includes("diff: -1 +2"), "diff summary should be shown"); + assertEquals(line.toolCallLabel, "Edit(/tmp/a.md)"); + assertEquals(line.toolStatus, "done"); + assertEquals(line.body, "ok"); assertEquals(line.diff?.map((row) => row.kind), [ "context", "remove", @@ -1242,13 +1325,13 @@ Deno.test("projectConsole renders snapshot entries and in-flight output", () => assertEquals(projection.status, "running"); assertEquals( projection.lines.map((line) => - `${line.kind}:${line.body}:${line.streaming}` + `${line.kind}:${line.toolCallLabel ? `${line.toolCallLabel}\n${line.body}` : line.body}:${line.streaming}` ), [ "user:seed user:false", "user:new user:false", "assistant:assistant reply:false", - "tool:Read — 1 file read\n /tmp/a.md:false", + "tool:Read(1 file)\n /tmp/a.md:false", "status:Compacting…:true", "in_flight:partial:true", ], @@ -1476,25 +1559,26 @@ Deno.test("projectConsole relativizes known tool path displays from snapshot cwd }, ]); - const bodies = projection.lines.filter((line) => line.kind === "tool").map(( - line, - ) => line.body); - assertEquals(bodies[0], "Read — 1 file read\n src/main.rs"); + const toolLines = projection.lines.filter((line) => line.kind === "tool"); + const bodies = toolLines.map((line) => line.body); + assertEquals(toolLines[0].toolCallLabel, "Read(1 file)"); + assertEquals(bodies[0], " src/main.rs"); assert( projection.lines[0].detail?.includes("from src/main.rs"), "Read summary detail path should be relative", ); assert( - bodies.some((body) => - body.includes("Write — out.txt") && body.includes("Wrote out.txt") + toolLines.some((line) => + line.toolCallLabel === "Write(out.txt)" && line.body.includes("Wrote out.txt") ), - "Write header and known result path should be relative", + "Write signature and known result path should be relative", ); assert( - bodies.some((body) => - body.includes("Edit — src/main.rs") && body.includes("Edited src/main.rs") + toolLines.some((line) => + line.toolCallLabel === "Edit(src/main.rs)" && + line.body.includes("Edited src/main.rs") ), - "Edit header and known result path should be relative", + "Edit signature and known result path should be relative", ); assert( bodies.some((body) => diff --git a/web/workspace/src/lib/workspace/console/model.ts b/web/workspace/src/lib/workspace/console/model.ts index 67a0f29e..d9e56953 100644 --- a/web/workspace/src/lib/workspace/console/model.ts +++ b/web/workspace/src/lib/workspace/console/model.ts @@ -11,6 +11,7 @@ import type { InternalWorkerSnapshot, Segment, } from "$lib/generated/protocol"; +import { stringify as stringifyYaml } from "yaml"; import { workspaceRoute } from "$lib/workspace/api/http"; import { applyRunActivityEvent, @@ -86,6 +87,9 @@ export type ConsoleLine = { kind: ConsoleLineKind; title: string; body: string; + expandedBody?: string; + toolCallLabel?: string; + toolStatus?: string; detail?: string; compaction?: ConsoleCompaction; diff?: ConsoleDiffLine[]; @@ -1376,7 +1380,10 @@ function refreshedToolLine(item: ConsoleLine): ConsoleLine { title: item.title.startsWith("Call · Tool result") ? item.title : `Call · ${toolCall.name}`, - body: renderToolCall(toolCall), + body: renderToolResponse(toolCall), + expandedBody: renderToolResponse(toolCall, true), + toolCallLabel: toolCallSignature(toolCall), + toolStatus: toolCallStatus(toolCall), detail: toolCallDetail(toolCall), diff: toolCall.name === "Edit" ? editDiff(toolCall) : undefined, streaming: !["done", "error"].includes(toolCall.state) && !commandTerminal, @@ -1384,7 +1391,7 @@ function refreshedToolLine(item: ConsoleLine): ConsoleLine { }; } -function renderToolCall(toolCall: ToolCallView): string { +function renderToolResponse(toolCall: ToolCallView, expanded = false): string { switch (toolCall.name) { case "Read": return renderReadTool(toolCall); @@ -1395,14 +1402,54 @@ function renderToolCall(toolCall: ToolCallView): string { case "Glob": return renderSearchTool(toolCall); case "Grep": - return renderGrepTool(toolCall); + return renderGrepTool(toolCall, expanded); case "Bash": - return renderBashTool(toolCall); + return renderBashTool(toolCall, expanded); default: - return renderDefaultTool(toolCall); + return renderDefaultTool(toolCall, expanded); } } +function toolCallSignature(toolCall: ToolCallView): string { + const args = parsedArgs(toolCall); + switch (toolCall.name) { + case "Read": + return `Read(${readPath(toolCall)})`; + case "Write": + case "Edit": { + const path = displayPath(stringField(args, "file_path") ?? "?", toolCall.cwd); + return `${toolCall.name}(${path})`; + } + case "Glob": + return `Glob(${stringField(args, "pattern") ?? genericCallArguments(toolCall)})`; + case "Grep": + return `Grep(${stringField(args, "pattern") ?? genericCallArguments(toolCall)})`; + case "Bash": { + const command = stringField(args, "command"); + return `Bash(${command ? `$ ${singleLine(command)}` : genericCallArguments(toolCall)})`; + } + default: + return `${toolCall.name}(${genericCallArguments(toolCall)})`; + } +} + +function genericCallArguments(toolCall: ToolCallView): string { + const raw = toolCall.arguments ?? toolCall.argsStream; + if (!raw.trim()) return ""; + const parsed = parseJson(raw); + if (parsed === undefined) return singleLine(raw); + const serialized = JSON.stringify(parsed) ?? "null"; + return isRecord(parsed) ? serialized.slice(1, -1) : serialized; +} + +function singleLine(value: string): string { + return value.replace(/\s+/g, " ").trim(); +} + +function toolCallStatus(toolCall: ToolCallView): string { + return toolCall.name === "Bash" ? commandStateSuffix(toolCall) : stateSuffix(toolCall.state); +} + function aggregateReadToolLines(lines: ConsoleLine[]): ConsoleLine[] { const result: ConsoleLine[] = []; let index = 0; @@ -1434,9 +1481,6 @@ function readAggregateLine(group: ConsoleLine[]): ConsoleLine { const paths = calls.map(readPath); const visiblePaths = inProgress ? paths.slice(-3) : paths; const body = compactLines([ - inProgress - ? `Read — reading (${count} file${plural(count)}…)` - : `Read — ${count} file${plural(count)} read`, visiblePaths.map((path) => ` ${path}`).join("\n"), inProgress && paths.length > visiblePaths.length ? ` … (${paths.length - visiblePaths.length} earlier)` @@ -1447,6 +1491,8 @@ function readAggregateLine(group: ConsoleLine[]): ConsoleLine { kind: "tool", title: "Call · Read", body, + toolCallLabel: `Read(${count} file${plural(count)})`, + toolStatus: hasError ? "failed" : inProgress ? "reading…" : "done", detail: calls.map(readDetail).join("\n\n"), eventId: group.at(-1)?.eventId, source: "event", @@ -1483,32 +1529,16 @@ function readDetail(toolCall: ToolCallView): string { ]); } -function renderReadTool(toolCall: ToolCallView): string { - return `Read — ${readPath(toolCall)} (${stateSuffix(toolCall.state)})`; +function renderReadTool(_toolCall: ToolCallView): string { + return ""; } function renderWriteTool(toolCall: ToolCallView): string { - const args = parsedArgs(toolCall); - const path = displayPath(stringField(args, "file_path") ?? "?", toolCall.cwd); - const content = stringField(args, "content"); - return compactLines([ - `Write — ${path} (${stateSuffix(toolCall.state)})`, - cappedSection(content, 5), - knownToolResultText(toolCall), - ]); + return knownToolResultText(toolCall) ?? ""; } function renderEditTool(toolCall: ToolCallView): string { - const args = parsedArgs(toolCall); - const path = displayPath(stringField(args, "file_path") ?? "?", toolCall.cwd); - const diff = editDiff(toolCall) ?? []; - const removes = diff.filter((line) => line.kind === "remove").length; - const adds = diff.filter((line) => line.kind === "add").length; - return compactLines([ - `Edit — ${path} (${stateSuffix(toolCall.state)})`, - diff.length > 0 ? `diff: -${removes} +${adds}` : undefined, - knownToolResultText(toolCall), - ]); + return knownToolResultText(toolCall) ?? ""; } function editDiff(toolCall: ToolCallView): ConsoleDiffLine[] | undefined { @@ -1591,52 +1621,20 @@ function lcsTable(oldLines: string[], newLines: string[]): number[][] { } function renderSearchTool(toolCall: ToolCallView): string { - const summary = toolCall.summary?.trim(); - return compactLines([ - `${toolCall.name} — ${toolHeaderSuffix(toolCall, summary)}`, - knownToolResultText(toolCall), - ]); + return knownToolResultText(toolCall) ?? ""; } -function renderGrepTool(toolCall: ToolCallView): string { - const summary = toolCall.summary?.trim(); - return compactLines([ - `Grep — ${toolHeaderSuffix(toolCall, summary)}`, - grepQueryText(toolCall), - cappedResultSection(knownToolResultText(toolCall), 5), - ]); +function renderGrepTool(toolCall: ToolCallView, expanded: boolean): string { + const result = knownToolResultText(toolCall); + return expanded ? result ?? "" : cappedResultSection(result, 5) ?? ""; } -function toolHeaderSuffix( - toolCall: ToolCallView, - summary?: string, -): string { - if (toolCall.state === "error") { - return "Failed"; +function renderBashTool(toolCall: ToolCallView, expanded: boolean): string { + if (["done", "error"].includes(toolCall.state)) { + const result = resultText(toolCall); + return expanded ? result ?? "" : cappedDisplaySection(result, 10) ?? ""; } - return summary ? firstLine(summary) : stateSuffix(toolCall.state); -} - -function grepQueryText(toolCall: ToolCallView): string | undefined { - const args = parsedArgs(toolCall); - const pattern = stringField(args, "pattern"); - if (pattern) { - return `query: ${pattern}`; - } - const renderedArgs = argsText(toolCall); - return renderedArgs ? `query:\n${renderedArgs}` : undefined; -} - -function renderBashTool(toolCall: ToolCallView): string { - const args = parsedArgs(toolCall); - const command = stringField(args, "command"); - return compactLines([ - `Bash — ${commandStateSuffix(toolCall)}`, - command ? `$ ${command}` : argsText(toolCall), - ["done", "error"].includes(toolCall.state) - ? cappedDisplaySection(resultText(toolCall), 10) - : renderLiveCommandOutput(toolCall.command), - ]); + return renderLiveCommandOutput(toolCall.command) ?? ""; } function commandStateSuffix(toolCall: ToolCallView): string { @@ -1690,12 +1688,9 @@ function renderLiveCommandOutput(command?: CommandSnapshot): string | undefined ]); } -function renderDefaultTool(toolCall: ToolCallView): string { - return compactLines([ - `${toolCall.name} — ${stateSuffix(toolCall.state)}`, - cappedDisplaySection(argsText(toolCall), 3), - cappedDisplaySection(resultText(toolCall), 3), - ]); +function renderDefaultTool(toolCall: ToolCallView, expanded: boolean): string { + const result = resultText(toolCall); + return expanded ? result ?? "" : cappedDisplaySection(result, 3) ?? ""; } function toolCallDetail(toolCall: ToolCallView): string { @@ -1713,10 +1708,30 @@ function toolCallDetail(toolCall: ToolCallView): string { } function resultText(toolCall: ToolCallView): string | undefined { - if (toolCall.output) { - return toolCall.output; + const text = toolCall.output || toolCall.summary; + return text ? formatJsonResponseAsYaml(text) : undefined; +} + +function formatJsonResponseAsYaml(text: string): string { + const trimmed = text.trim(); + if ( + !( + (trimmed.startsWith("{") && trimmed.endsWith("}")) || + (trimmed.startsWith("[") && trimmed.endsWith("]")) + ) + ) { + return text; + } + + try { + const parsed: unknown = JSON.parse(trimmed); + if (parsed === null || typeof parsed !== "object") { + return text; + } + return stringifyYaml(parsed).trimEnd(); + } catch { + return text; } - return toolCall.summary; } function knownToolResultText(toolCall: ToolCallView): string | undefined { @@ -1803,7 +1818,7 @@ function argsText(toolCall: ToolCallView): string { return ""; } const parsed = parseJson(raw); - return parsed === undefined ? raw : jsonPreview(parsed); + return parsed === undefined ? raw : stringifyYaml(parsed).trimEnd(); } function parsedArgs( @@ -1830,21 +1845,6 @@ function compactLines(lines: Array): string { return lines.filter((line): line is string => Boolean(line)).join("\n"); } -function cappedSection( - value: string | undefined, - cap: number, -): string | undefined { - if (!value) { - return undefined; - } - const lines = value.split(/\r?\n/); - const shown = lines.slice(0, cap); - if (lines.length > cap) { - shown.push(`… +${lines.length - cap} more lines`); - } - return shown.join("\n"); -} - function cappedDisplaySection( value: string | undefined, maxLines: number, diff --git a/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts b/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts index 1fcb2b37..785d9128 100644 --- a/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts +++ b/web/workspace/src/lib/workspace/console/worker-console.ui.test.ts @@ -403,7 +403,7 @@ Deno.test("Worker Console renders markdown only for message rows", async () => { consoleLine.includes("item.kind === 'tool'") && consoleLine.includes("{#if isBashTool(item)}") && consoleLine.includes( - "", + "", ) && consoleLine.includes( ".console-line.tool-bash .console-plain-text", @@ -420,6 +420,30 @@ Deno.test("Worker Console renders markdown only for message rows", async () => { ); }); +Deno.test("Worker Console expands uncapped tool body from the hover detail action", async () => { + const consoleLine = await Deno.readTextFile( + new URL("./ConsoleLineItem.svelte", import.meta.url), + ); + + assert( + consoleLine.includes( + "return detailOpen ? (line.expandedBody ?? line.body) : line.body", + ) && + consoleLine.includes("line.toolCallLabel ?? line.toolCall?.name") && + consoleLine.includes('class={`tool-status') && + consoleLine.includes('class="tool-detail-button"') && + consoleLine.includes("aria-expanded={detailOpen}") && + consoleLine.includes("detailOpen = !detailOpen") && + consoleLine.includes("item.detail && detailOpen") && + consoleLine.includes('role="region"') && + consoleLine.includes(".console-line:hover .tool-detail-button") && + consoleLine.includes(".tool-detail-button:focus-visible") && + consoleLine.includes("@media (hover: none)") && + !consoleLine.includes('
'), + "Normal tool display should keep its preview while detail reveals the uncapped body and existing metadata", + ); +}); + Deno.test("Worker Console renders Edit diffs without preformatted template gaps", async () => { const consoleLine = await Deno.readTextFile( new URL("./ConsoleLineItem.svelte", import.meta.url), diff --git a/web/workspace/src/lib/workspace/sidebar/WorkersNavSection.svelte b/web/workspace/src/lib/workspace/sidebar/WorkersNavSection.svelte index b2fffcd6..bb6a0743 100644 --- a/web/workspace/src/lib/workspace/sidebar/WorkersNavSection.svelte +++ b/web/workspace/src/lib/workspace/sidebar/WorkersNavSection.svelte @@ -5,7 +5,7 @@ workspaceWorkersStore, type SidebarWorker, } from './worker-subscription'; - import { canShowWorkerInSidebar } from './workers'; + import { canShowWorkerInSidebar, sidebarWorkerActivity } from './workers'; const COLLAPSED_WORKER_COUNT = 6; @@ -69,6 +69,7 @@