diff --git a/crates/agen/README.md b/crates/agen/README.md index 46bd05d0..f9dba8c9 100644 --- a/crates/agen/README.md +++ b/crates/agen/README.md @@ -21,20 +21,22 @@ 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; +use agen::{Engine, EngineError, History}; use agen::llm_client::LlmClient; -async fn conversation(client: C) { +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 _exit = 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 7779ea03..5f9d9874 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, EngineRunExit, StopReason}; +use agen::{Engine, EngineResult, History}; 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 = History::new(); println!("🚀 Starting Engine..."); println!("💡 Will cancel after 2 seconds\n"); @@ -45,13 +46,15 @@ async fn main() -> Result<(), Box> { println!("📡 Sending request to LLM..."); - let output = engine.run("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") + match engine.run(&mut history, "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); } EngineRunExit::Interrupted(reason) => println!("❌ Task interrupted: {reason:?}"), } diff --git a/crates/agen/examples/engine_cli.rs b/crates/agen/examples/engine_cli.rs index a275ba9a..1c9abfd2 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, EngineRunExit, StopReason, + Engine, History, interceptor::{Interceptor, PostToolAction, ToolResultInfo}, llm_client::{ LlmClient, @@ -474,11 +474,16 @@ async fn main() -> Result<(), Box> { engine.set_interceptor(ToolResultPrinterPolicy::new(tool_call_names)); + let mut history = History::new(); + // One-shot mode if let Some(prompt) = args.prompt { - let output = engine.run(&prompt).await; - if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = output.result { - eprintln!("\n❌ Error: {error}"); + match engine.run(&mut history, &prompt).await { + Ok(_) => {} + Err(e) => { + eprintln!("\n❌ Error: {}", e); + std::process::exit(1); + } } return Ok(()); @@ -497,8 +502,13 @@ async fn main() -> Result<(), Box> { return Ok(()); } - let output = engine.run(first_input).await; - let mut locked = output.engine; + let mut locked = match engine.run(&mut history, first_input).await { + Ok(out) => out.engine, + Err(e) => { + eprintln!("\n❌ Error: {}", e); + return Ok(()); + } + }; loop { print!("\n👤 You: "); @@ -517,8 +527,11 @@ async fn main() -> Result<(), Box> { break; } - if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = locked.run(input).await { - eprintln!("\n❌ Error: {error}"); + match locked.run(&mut history, input).await { + Ok(_) => {} + Err(e) => { + eprintln!("\n❌ Error: {}", e); + } } } diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 60e9b0d4..4139d0b4 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -7,7 +7,7 @@ use tokio::sync::mpsc; use tracing::{debug, info, trace, warn}; use crate::{ - Item, + History, HistoryEntry, Item, callback::{ ClosureMetaHandler, ClosureTextBlockHandler, ClosureThinkingBlockHandler, ClosureToolUseBlockHandler, TextBlockScope, ThinkingBlockScope, ToolUseBlockScope, @@ -117,9 +117,9 @@ impl From> for EngineRunExit { /// 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: EngineRunExit, } @@ -139,29 +139,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 -/// let _exit = 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)] @@ -181,7 +183,7 @@ enum StreamCompletion { Interrupted { reason: String }, } -pub struct Engine { +pub struct Engine { /// LLM client client: C, /// Retry policy for opening an LLM response stream. @@ -201,8 +203,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. @@ -290,10 +290,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 { +impl Engine { fn start_logical_run(&mut self) { self.active_run_turn_count = Some(0); } @@ -559,11 +559,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(()) } @@ -670,9 +674,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 @@ -929,20 +933,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, @@ -1142,19 +1146,26 @@ impl Engine { } /// Internal turn execution logic - async fn run_turn_loop(&mut self) -> Result { + 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); } } @@ -1199,13 +1210,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 @@ -1273,7 +1284,7 @@ impl Engine { 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 { @@ -1289,7 +1300,7 @@ impl Engine { 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 => {} @@ -1348,7 +1359,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, @@ -1376,15 +1387,16 @@ 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 => { return Ok(EngineResult::Finished); } TurnEndAction::ContinueWithMessages(additional) => { - self.append_history_items(additional)?; + self.append_history_items(history, additional, annotate)?; continue; } TurnEndAction::Pause => { @@ -1393,7 +1405,10 @@ impl Engine { } } - 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); } } @@ -1639,6 +1654,8 @@ 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 { @@ -1655,7 +1672,7 @@ impl Engine { result.attachments, ) }); - self.append_history_items(items)?; + self.append_history_items(history, items, annotate)?; Ok(None) } Err(err) => Err(err), @@ -1663,9 +1680,9 @@ impl Engine { } } -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(); @@ -1687,7 +1704,6 @@ 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, @@ -1847,36 +1863,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) @@ -1895,16 +1913,21 @@ impl Engine { 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(self, user_input: impl Into) -> EngineRunOutput { - let mut locked = self.lock(); - let result = locked.run(user_input).await; + /// 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, + 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, @@ -1914,9 +1937,13 @@ impl Engine { /// Resume from Paused, consuming self and transitioning to Locked. /// /// Used after `unlock()` → edit → resume. - pub async fn resume(self) -> EngineRunOutput { - let mut locked = self.lock(); - let result = locked.resume().await; + 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, @@ -1936,9 +1963,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, @@ -1949,7 +1976,6 @@ 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, @@ -1983,19 +2009,73 @@ 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(&mut self, user_input: impl Into) -> EngineRunExit { - self.run_result(user_input.into()).await.into() + 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(&mut self, user_input: String) -> Result { + 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; - // 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) => { @@ -2006,27 +2086,35 @@ 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 = self.run_turn_loop(history, annotate).await; 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) -> EngineRunExit { - self.resume_result().await.into() + /// 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(&mut self) -> Result { + 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 = self.run_turn_loop(history, annotate).await; let result = self.finalize_interruption(result).await; self.finish_logical_run(&result); result @@ -2041,7 +2129,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, @@ -2052,7 +2140,6 @@ 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, @@ -2086,6 +2173,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), 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 0e4d0003..871af33b 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; @@ -24,6 +25,7 @@ pub use engine::{ 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}; 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/tests/annotated_history_test.rs b/crates/agen/tests/annotated_history_test.rs new file mode 100644 index 00000000..bb0872e9 --- /dev/null +++ b/crates/agen/tests/annotated_history_test.rs @@ -0,0 +1,85 @@ +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 + .unwrap(); + + assert!(matches!(output.result, agen::EngineResult::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 2e84b145..43d91053 100644 --- a/crates/agen/tests/callback_test.rs +++ b/crates/agen/tests/callback_test.rs @@ -8,11 +8,11 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; -use agen::Engine; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent as ClientStatusEvent}; use agen::llm_client::retry::RetryPolicy; use agen::llm_client::{ClientError, LlmClient, Request, ResponseStream}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; +use agen::{Engine, History}; use async_trait::async_trait; use common::MockLlmClient; @@ -58,6 +58,7 @@ async fn test_callback_llm_retry_event() { max_attempts: 2, total_timeout: Duration::from_secs(1), }); + let mut history: History = History::new(); let notices = Arc::new(Mutex::new(Vec::new())); let sink = notices.clone(); @@ -65,11 +66,8 @@ async fn test_callback_llm_retry_event() { sink.lock().unwrap().push((llm_call, notice.clone())); }); - let result = engine.run("retry once").await; - assert!( - matches!(result.result, agen::EngineRunExit::Finished), - "engine should succeed after one retry" - ); + let result = engine.run(&mut history, "retry once").await; + assert!(result.is_ok(), "engine should succeed after one retry"); let notices = notices.lock().unwrap(); assert_eq!(notices.len(), 1); @@ -94,6 +92,7 @@ async fn test_callback_text_block_events() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history: History = History::new(); let text_deltas = Arc::new(Mutex::new(Vec::new())); let text_completes = Arc::new(Mutex::new(Vec::new())); @@ -111,12 +110,9 @@ async fn test_callback_text_block_events() { }); }); - // Mutable::run consumes self, returns (Locked, EngineRunExit) - let result = engine.run("Greet me").await; - assert!( - matches!(result.result, agen::EngineRunExit::Finished), - "Engine should complete" - ); + // Mutable::run consumes self, returns (Locked, EngineResult) + let result = engine.run(&mut history, "Greet me").await; + assert!(result.is_ok(), "Engine should complete"); let deltas = text_deltas.lock().unwrap(); assert_eq!(deltas.len(), 2); @@ -143,6 +139,7 @@ async fn test_callback_tool_call_complete() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history: History = History::new(); let tool_starts = Arc::new(Mutex::new(Vec::<(String, String)>::new())); let tool_completes = Arc::new(Mutex::new(Vec::new())); @@ -160,8 +157,8 @@ async fn test_callback_tool_call_complete() { }); }); - // Mutable::run consumes self, returns (Locked, EngineRunExit) - let _ = engine.run("Weather please").await; + // Mutable::run consumes self, returns (Locked, EngineResult) + let _ = engine.run(&mut history, "Weather please").await; let starts = tool_starts.lock().unwrap(); assert_eq!(starts.len(), 1); @@ -189,6 +186,7 @@ async fn test_callback_turn_events() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history: History = History::new(); let turn_starts = Arc::new(Mutex::new(Vec::new())); let turn_ends = Arc::new(Mutex::new(Vec::new())); @@ -203,9 +201,9 @@ async fn test_callback_turn_events() { ends.lock().unwrap().push(turn); }); - // Mutable::run consumes self, returns (Locked, EngineRunExit) - let result = engine.run("Do something").await; - assert!(matches!(result.result, agen::EngineRunExit::Finished)); + // Mutable::run consumes self, returns (Locked, EngineResult) + let result = engine.run(&mut history, "Do something").await; + assert!(result.is_ok()); let starts = turn_starts.lock().unwrap(); let ends = turn_ends.lock().unwrap(); @@ -260,6 +258,7 @@ async fn test_callback_tool_result_events() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history: History = History::new(); engine.register_tool(fixed_tool( "fixed", @@ -282,7 +281,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); @@ -336,6 +335,7 @@ async fn test_callback_tool_result_error_path() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history: History = History::new(); engine.register_tool(erroring_tool("erroring", "boom")); @@ -351,7 +351,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); @@ -380,6 +380,7 @@ async fn test_callback_usage_events() { let client = MockLlmClient::new(events); let mut engine = Engine::new(client); + let mut history: History = History::new(); let usage_events = Arc::new(Mutex::new(Vec::new())); @@ -388,8 +389,8 @@ async fn test_callback_usage_events() { usages.lock().unwrap().push(event.clone()); }); - // Mutable::run consumes self, returns (Locked, EngineRunExit) - let _ = engine.run("Hello").await; + // Mutable::run consumes self, returns (Locked, EngineResult) + let _ = engine.run(&mut history, "Hello").await; let usages = usage_events.lock().unwrap(); assert_eq!(usages.len(), 1); diff --git a/crates/agen/tests/engine_fixtures.rs b/crates/agen/tests/engine_fixtures.rs index c6f1f4f7..c5471beb 100644 --- a/crates/agen/tests/engine_fixtures.rs +++ b/crates/agen/tests/engine_fixtures.rs @@ -9,8 +9,8 @@ use std::path::Path; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use agen::Engine; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; +use agen::{Engine, History}; use async_trait::async_trait; use common::MockLlmClient; @@ -134,9 +134,10 @@ async fn test_engine_simple_text_response() { let client = MockLlmClient::from_fixture(&fixture_path).unwrap(); let engine = Engine::new(client); + let mut history: History = 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!( matches!(result.result, agen::EngineRunExit::Finished), @@ -159,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: History = History::new(); // Register tool let weather_tool = MockWeatherTool::new(); @@ -166,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 @@ -198,9 +202,10 @@ async fn test_engine_with_programmatic_events() { let client = MockLlmClient::new(events); let engine = Engine::new(client); + let mut history: History = 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!( matches!(result.result, agen::EngineRunExit::Finished), diff --git a/crates/agen/tests/engine_state_test.rs b/crates/agen/tests/engine_state_test.rs index 5b7c082e..c72f136d 100644 --- a/crates/agen/tests/engine_state_test.rs +++ b/crates/agen/tests/engine_state_test.rs @@ -15,7 +15,7 @@ use agen::interceptor::{ use agen::llm_client::ClientError; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; -use agen::{Engine, EngineError, EngineRunExit, StopReason}; +use agen::{Engine, EngineError, EngineResult, History}; use async_trait::async_trait; use common::MockLlmClient; @@ -75,36 +75,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 @@ -112,9 +113,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. @@ -124,6 +126,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()); @@ -132,18 +135,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"] @@ -218,6 +224,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() { @@ -227,8 +234,8 @@ async fn history_append_failure_stops_before_tool_execution() { } }); - let mut engine = engine.lock(); - let exit = engine.run("use the tool").await; + let mut engine = engine.lock(&history); + let error = engine.run(&mut history, "use the tool").await.unwrap_err(); assert!(matches!( exit, @@ -236,8 +243,8 @@ async fn history_append_failure_stops_before_tool_execution() { 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")); } // ============================================================================= @@ -249,21 +256,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); } @@ -272,21 +280,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()); } // ============================================================================= @@ -307,20 +316,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(()) } @@ -351,35 +360,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!(matches!(result1, EngineRunExit::Finished)); - assert_eq!(locked_engine.history().len(), 2); // user + assistant + let result1 = locked_engine.run(&mut history, "Hello!").await; + assert!(result1.is_ok()); + assert_eq!(history.len(), 2); // user + assistant // Turn 2 - let result2 = locked_engine.run("Can you help me?").await; - assert!(matches!(result2, EngineRunExit::Finished)); - assert_eq!(locked_engine.history().len(), 4); // 2 * (user + assistant) + let result2 = locked_engine.run(&mut history, "Can you help me?").await; + assert!(result2.is_ok()); + 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 @@ -405,26 +415,36 @@ 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; + locked_engine + .run(&mut history, "New message") + .await + .unwrap(); // 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 } @@ -451,18 +471,19 @@ 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; + engine.run(&mut history, "Second").await?; assert_eq!(engine.turn_count(), 2); assert_eq!(engine.llm_call_count(), 2); @@ -482,28 +503,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); } @@ -546,19 +568,23 @@ 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; + let mut locked = engine.lock(&history); + locked.run(&mut history, "first").await.expect("first run"); 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; + let mut relocked = unlocked.lock(&history); + relocked + .run(&mut history, "second") + .await + .expect("second run"); 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"); @@ -573,8 +599,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(); @@ -589,14 +616,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")); } @@ -660,17 +688,21 @@ impl Interceptor for ContinueTurnOnce { async fn max_turns_is_scoped_to_each_fresh_run() { let responses = vec![completed_text_events(), completed_text_events()]; let mut engine = Engine::new(MockLlmClient::with_responses(responses)); + let mut history: History = History::new(); engine.set_max_turns(Some(1)); - let mut engine = engine.lock(); + let mut engine = engine.lock(&history); - assert!(matches!(engine.run("first").await, EngineRunExit::Finished)); + assert_eq!( + engine.run(&mut history, "first").await.unwrap(), + EngineResult::Finished + ); assert_eq!(engine.turn_count(), 1); assert_eq!(engine.active_run_turn_count(), None); - assert!(matches!( - engine.run("second").await, - EngineRunExit::Finished - )); + assert_eq!( + engine.run(&mut history, "second").await.unwrap(), + EngineResult::Finished + ); assert_eq!(engine.turn_count(), 2); assert_eq!(engine.active_run_turn_count(), None); } @@ -678,17 +710,24 @@ async fn max_turns_is_scoped_to_each_fresh_run() { #[tokio::test] async fn yielded_resume_keeps_the_same_unspent_turn_budget() { let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); + let mut history: History = History::new(); 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!(matches!(engine.run("start").await, EngineRunExit::Yielded)); + assert_eq!( + engine.run(&mut history, "start").await.unwrap(), + EngineResult::Yielded + ); assert_eq!(engine.turn_count(), 0); assert_eq!(engine.active_run_turn_count(), Some(0)); - assert!(matches!(engine.resume().await, EngineRunExit::Finished)); + assert_eq!( + engine.resume(&mut history).await.unwrap(), + EngineResult::Finished + ); assert_eq!(engine.turn_count(), 1); assert_eq!(engine.active_run_turn_count(), None); } @@ -705,22 +744,26 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() { ]; let tool = CountingTool::new("count_tool"); let mut engine = Engine::new(MockLlmClient::new(events)); + let mut history: History = History::new(); engine.set_max_turns(Some(1)); engine.register_tool(tool.definition()); engine.set_interceptor(PauseToolOnce { calls: AtomicUsize::new(0), }); - let mut engine = engine.lock(); + let mut engine = engine.lock(&history); - assert!(matches!(engine.run("call it").await, EngineRunExit::Paused)); + assert_eq!( + engine.run(&mut history, "call it").await.unwrap(), + EngineResult::Paused + ); assert_eq!(engine.turn_count(), 1); assert_eq!(engine.active_run_turn_count(), Some(1)); assert_eq!(tool.call_count(), 0); - assert!(matches!( - engine.resume().await, - EngineRunExit::Interrupted(StopReason::LimitReached) - )); + assert_eq!( + engine.resume(&mut history).await.unwrap(), + EngineResult::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"); @@ -739,20 +782,24 @@ async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() { let client = MockLlmClient::with_responses(vec![tool_events, completed_text_events()]); let tool = CountingTool::new("count_tool"); let mut engine = Engine::new(client); + let mut history: History = History::new(); engine.set_max_turns(Some(1)); engine.register_tool(tool.definition()); engine.set_interceptor(PauseToolOnce { calls: AtomicUsize::new(0), }); - let mut engine = engine.lock(); + let mut engine = engine.lock(&history); - assert!(matches!(engine.run("pause").await, EngineRunExit::Paused)); + assert_eq!( + engine.run(&mut history, "pause").await.unwrap(), + EngineResult::Paused + ); assert_eq!(engine.active_run_turn_count(), Some(1)); - assert!(matches!( - engine.run("replace").await, - EngineRunExit::Finished - )); + assert_eq!( + engine.run(&mut history, "replace").await.unwrap(), + EngineResult::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"); @@ -761,16 +808,17 @@ 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 engine = Engine::new(MockLlmClient::new(completed_text_events())); + let mut history: History = History::new(); 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!(matches!( - engine.run("start").await, - EngineRunExit::Interrupted(StopReason::LimitReached) - )); + assert_eq!( + engine.run(&mut history, "start").await.unwrap(), + EngineResult::LimitReached + ); assert_eq!(engine.turn_count(), 1); assert_eq!(engine.llm_call_count(), 1); assert_eq!(engine.active_run_turn_count(), None); @@ -779,15 +827,16 @@ 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 engine = Engine::new(MockLlmClient::new(completed_text_events())); + let mut history: History = History::new(); engine.set_max_turns(Some(1)); engine.set_turn_count(7); engine.set_active_run_turn_count(Some(1)); - let mut engine = engine.lock(); + let mut engine = engine.lock(&history); - assert!(matches!( - engine.resume().await, - EngineRunExit::Interrupted(StopReason::LimitReached) - )); + assert_eq!( + engine.resume(&mut history).await.unwrap(), + EngineResult::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 3ec4647f..42134f81 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -6,12 +6,12 @@ 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, }; +use agen::{Engine, History}; use async_trait::async_trait; mod common; @@ -145,6 +145,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); @@ -158,8 +159,8 @@ async fn test_parallel_tool_execution() { engine.register_tool(tool3.definition()); let start = Instant::now(); - // Mutable::run consumes self, returns (Locked, EngineRunExit) - let _result = engine.run("Run all tools").await; + // Mutable::run consumes self, returns (Locked, EngineResult) + let _result = engine.run(&mut history, "Run all tools").await; let elapsed = start.elapsed(); // Verify all tools were called @@ -205,13 +206,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 +258,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 +301,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 +348,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 +395,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); @@ -415,8 +422,8 @@ async fn test_before_tool_call_skip() { engine.set_interceptor(BlockingPolicy); - // Mutable::run consumes self, returns (Locked, EngineRunExit) - let _result = engine.run("Test hook").await; + // Mutable::run consumes self, returns (Locked, EngineResult) + let _result = engine.run(&mut history, "Test hook").await; // allowed_tool is called, but blocked_tool is not assert_eq!( @@ -457,6 +464,7 @@ async fn test_post_tool_call_modification() { ]); let mut engine = Engine::new(client); + let mut history: History = History::new(); #[derive(Clone)] struct SimpleTool; @@ -502,8 +510,8 @@ async fn test_post_tool_call_modification() { modified_content: modified_content.clone(), }); - // Mutable::run consumes self, returns (Locked, EngineRunExit) - let result = engine.run("Test modification").await; + // Mutable::run consumes self, returns (Locked, EngineResult) + let result = engine.run(&mut history, "Test modification").await; assert!( matches!(result.result, agen::EngineRunExit::Finished), @@ -543,6 +551,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()); @@ -561,10 +570,13 @@ async fn test_before_tool_call_synthetic_result_committed() { engine.set_interceptor(SyntheticPolicy); - let result = engine.run("Test synthetic result").await; + let _result = engine + .run(&mut history, "Test synthetic result") + .await + .unwrap(); 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, diff --git a/crates/agen/tests/reasoning_round_trip_test.rs b/crates/agen/tests/reasoning_round_trip_test.rs index f374e85a..691466c6 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; - let engine = out.engine; + let mut history: History = History::new(); + let _out = engine.run(&mut history, "question?").await.expect("run ok"); - 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; - let engine = out.engine; + let mut history: History = History::new(); + let _out = engine.run(&mut history, "q").await.expect("run ok"); - 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; - let engine = out.engine; + let mut history: History = History::new(); + let _out = engine.run(&mut history, "q").await.expect("run ok"); - 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; + let _ = engine.run(&mut history, "follow up").await.expect("run ok"); 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/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/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/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 2116370f..484047b5 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::EngineRunExit) { +) -> (TestWorker, agen::EngineResult) { // 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,13 +147,14 @@ 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(); @@ -178,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(); @@ -222,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( @@ -230,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(); @@ -261,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( @@ -269,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(); @@ -295,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); @@ -304,7 +338,7 @@ async fn session_resume_after_pause() { SegmentStartState { system_prompt: worker.get_system_prompt(), config: worker.request_config(), - history: worker.history(), + history: &worker.history(), }, ) .unwrap(); @@ -335,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( @@ -343,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(); @@ -356,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(); @@ -377,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(); @@ -422,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(); @@ -455,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(); @@ -488,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(); @@ -540,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/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index fec60a5d..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) } } @@ -2507,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, @@ -3243,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 ccc2acd1..b14e45df 100644 --- a/crates/worker/examples/worker_cli.rs +++ b/crates/worker/examples/worker_cli.rs @@ -71,7 +71,7 @@ async fn main() -> Result<(), Box> { } // 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 638adadc..82ab1a08 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1780,7 +1780,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() { @@ -1806,7 +1806,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) { @@ -1854,7 +1854,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(); 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/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/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/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 e0464198..a8d21e49 100644 --- a/crates/worker/src/internal_worker.rs +++ b/crates/worker/src/internal_worker.rs @@ -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()); @@ -536,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()); @@ -633,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 { 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 76943f1d..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; @@ -42,6 +43,10 @@ pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTem 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/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 11b86f30..21a9371a 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -1252,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")))) })); @@ -1274,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( diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 9207a232..d6d32819 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -5,18 +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, EngineRunExit, StopReason, ToolOutputLimits, UsageRecord, + Engine, EngineError, EngineResult, EngineRunExit, History, HistoryEntry, Item, StopReason, + 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, @@ -25,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, @@ -805,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 @@ -830,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); }); @@ -858,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)) } } @@ -899,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 @@ -906,9 +1049,10 @@ pub const WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN: &str = "worker.input-submiss pub struct Worker { manifest: WorkerManifest, /// Always `Some` outside of `run()`/`resume()`. - engine: Option>, - /// Worker-owned recovery marker. Agen exposes only the typed run exit and - /// never persists or restores Worker lifecycle state. + 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 @@ -1008,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 @@ -1107,11 +1255,13 @@ 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, @@ -1141,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, @@ -1159,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 @@ -1205,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() @@ -1259,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 — @@ -1281,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 @@ -1295,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, @@ -1314,6 +1453,7 @@ 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, @@ -1343,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, @@ -1552,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) } @@ -1630,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") } @@ -1638,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, @@ -1789,9 +1963,14 @@ 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.last_run_interrupted = state.last_run_interrupted; @@ -1863,9 +2042,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. @@ -2169,7 +2357,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, @@ -2177,6 +2365,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( @@ -2299,7 +2488,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 } @@ -2337,7 +2529,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?; @@ -2362,7 +2557,7 @@ 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(), @@ -2390,8 +2585,9 @@ impl Worker { 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) } @@ -2399,7 +2595,8 @@ impl Worker { &mut self, snapshot: EmptyTurnRollbackSnapshot, ) -> Result<(), StoreError> { - self.engine_mut().truncate_history(snapshot.history_len); + 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); @@ -2430,7 +2627,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 { @@ -2439,7 +2643,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( @@ -2514,16 +2718,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. @@ -2537,7 +2740,10 @@ 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 } @@ -2545,8 +2751,11 @@ impl Worker { &mut self, input: Vec, mut input_extensions: Vec, - ) -> Result { - let (input, pending_flow_state) = self.prepare_flow_input(input)?; + ) -> Result + where + St: Clone + 'static, + { + 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}")) @@ -2579,13 +2788,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 @@ -2599,7 +2813,6 @@ impl Worker { // 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 @@ -2607,13 +2820,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)?; @@ -2698,7 +2939,10 @@ impl Worker { /// 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> { + fn apply_interrupt_prep(&mut self) -> Result<(), WorkerError> + where + St: Clone + 'static, + { let tool_result_summary = self .prompts() .load_full() @@ -2710,24 +2954,57 @@ impl Worker { .interrupt_system_note() .map_err(WorkerError::from)?; + let history_items = self.history(); let closures = crate::interrupt_prep::orphan_tool_result_closures( - self.engine().history(), + &history_items, &tool_result_summary, ); if !closures.is_empty() { - self.engine_mut().append_history(closures)?; + let mut annotate = history_annotator( + self.log_writer_handle(), + Vec::new(), + 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(), closures, &mut annotate)?; + session.note_mutation(); } 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); + let mut annotate = history_annotator( + self.log_writer_handle(), + vec![interrupt_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(interrupt_entry.item), + &mut annotate, + )?; + session.note_mutation(); Ok(()) } @@ -2751,37 +3028,40 @@ impl Worker { 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. @@ -2795,7 +3075,10 @@ impl Worker { pub async fn run_for_notification( &mut self, kind: protocol::InvokeKind, - ) -> Result { + ) -> Result + where + St: Clone + 'static, + { debug_assert!( matches!( kind, @@ -2817,27 +3100,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 } @@ -2856,12 +3155,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, }; @@ -2886,12 +3191,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(), @@ -2935,7 +3246,10 @@ impl Worker { &mut self, result: EngineRunExit, history_before: usize, - ) -> Result { + ) -> Result + where + St: Clone + 'static, + { self.persist_turn(history_before, &result).await?; if matches!(result, EngineRunExit::Yielded) { @@ -3024,7 +3338,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. @@ -3196,9 +3513,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 { @@ -3403,11 +3720,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 ( @@ -3535,9 +3859,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(), ), ), ), @@ -3814,6 +4138,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 @@ -3825,12 +4179,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, @@ -3894,7 +4251,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 @@ -4134,12 +4491,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( @@ -4232,9 +4584,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 @@ -4284,9 +4635,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( @@ -4770,7 +5121,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)); @@ -4780,6 +5134,7 @@ where let mut worker = Self { manifest, engine: Some(worker), + session: WorkerSession::new(session_id, Vec::new()), last_run_interrupted: false, store, worker_metadata_writer, @@ -4809,6 +5164,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, @@ -4851,7 +5207,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); @@ -4859,6 +5218,7 @@ 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, @@ -4888,6 +5248,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, @@ -4963,7 +5324,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)); @@ -4973,6 +5337,7 @@ where let mut worker = Self { manifest, engine: Some(worker), + session: WorkerSession::new(session_id, Vec::new()), last_run_interrupted: false, store, worker_metadata_writer, @@ -5002,6 +5367,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, @@ -5252,7 +5618,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 { @@ -5262,15 +5631,21 @@ 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_active_run_turn_count(state.active_run_turn_count); @@ -5287,6 +5662,7 @@ 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, @@ -5320,6 +5696,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, @@ -5417,7 +5794,10 @@ 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.set_tool_output_limits(Some(ToolOutputLimits { @@ -7087,7 +7467,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, @@ -7130,7 +7510,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, @@ -7169,7 +7549,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, @@ -7190,7 +7570,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()), @@ -7224,7 +7604,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()), @@ -7259,7 +7639,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(), @@ -7272,11 +7652,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 @@ -7284,11 +7666,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); @@ -7314,7 +7704,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()), @@ -7347,7 +7737,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, @@ -7473,11 +7863,8 @@ mod build_summary_prompt_tests { .len(), expected_truncate_entries ); - assert_eq!(worker.engine().history().len(), 1); - assert_eq!( - worker.engine().history()[0].as_text().unwrap(), - "first message" - ); + assert_eq!(worker.history().len(), 1); + assert_eq!(worker.history()[0].as_text().unwrap(), "first message"); } #[tokio::test] @@ -7506,7 +7893,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, @@ -7517,13 +7904,11 @@ 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-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!( @@ -7547,8 +7932,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" ) @@ -7559,8 +7947,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 ) @@ -7582,7 +7973,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, @@ -7600,7 +7991,7 @@ mod build_summary_prompt_tests { item: dangling_call.clone().into(), }) .unwrap(); - worker.engine_mut().set_history(vec![dangling_call]); + worker.set_history_for_test(vec![dangling_call]); worker.last_run_interrupted = true; worker @@ -7608,7 +7999,7 @@ mod build_summary_prompt_tests { .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" @@ -7658,7 +8049,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()), @@ -7718,7 +8109,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, @@ -7911,7 +8302,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()), @@ -7949,8 +8340,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") @@ -7981,7 +8375,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()), @@ -7997,7 +8391,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(), @@ -8056,7 +8450,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..04836dbe 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, @@ -804,10 +826,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 +839,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 +1113,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 +1344,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 +1487,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 +1573,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!( @@ -2373,7 +2391,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,