Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

feat: add provenance-aware worker history

Browse Source
This commit is contained in:
Hare
2026-08-27 12:51:01 +09:00
parent 0496cd907b
commit 374449e663
46 changed files with 2505 additions and 605 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/README.md
+5 -4
View File
@@ -21,20 +21,21 @@ agen = { version = "0.2.1", features = ["codex"] }
## Quick start ## 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 ```no_run
use agen::{Engine, EngineError}; use agen::{Engine, EngineError, History};
use agen::llm_client::LlmClient; use agen::llm_client::LlmClient;
async fn conversation<C: LlmClient>(client: C) -> Result<(), EngineError> { async fn conversation<C: LlmClient>(client: C) -> Result<(), EngineError> {
let mut history = History::new();
let output = Engine::new(client) let output = Engine::new(client)
.system_prompt("You are a concise assistant.") .system_prompt("You are a concise assistant.")
.run("Explain typed state in one sentence.") .run(&mut history, "Explain typed state in one sentence.")
.await?; .await?;
let mut engine = output.engine; let mut engine = output.engine;
let _result = engine.run("Give a Rust example.").await?; let _result = engine.run(&mut history, "Give a Rust example.").await?;
Ok(()) Ok(())
} }
``` ```
crates/agen/examples/engine_cancel_demo.rs
+3 -2
View File
@@ -4,7 +4,7 @@
use agen::llm_client::scheme::{Scheme, anthropic::AnthropicScheme}; use agen::llm_client::scheme::{Scheme, anthropic::AnthropicScheme};
use agen::llm_client::transport::{HttpTransport, ResolvedAuth}; use agen::llm_client::transport::{HttpTransport, ResolvedAuth};
use agen::{Engine, EngineResult}; use agen::{Engine, EngineResult, History};
use std::time::Duration; use std::time::Duration;
#[tokio::main] #[tokio::main]
@@ -29,6 +29,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let base_url = scheme.default_base_url().to_string(); let base_url = scheme.default_base_url().to_string();
let client = HttpTransport::new(scheme, model, base_url, ResolvedAuth::ApiKey(api_key), cap); let client = HttpTransport::new(scheme, model, base_url, ResolvedAuth::ApiKey(api_key), cap);
let engine = Engine::new(client); let engine = Engine::new(client);
let mut history = History::new();
println!("🚀 Starting Engine..."); println!("🚀 Starting Engine...");
println!("💡 Will cancel after 2 seconds\n"); println!("💡 Will cancel after 2 seconds\n");
@@ -45,7 +46,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
println!("📡 Sending request to LLM..."); println!("📡 Sending request to LLM...");
match engine.run("Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await { 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 { Ok(out) => match out.result {
EngineResult::Finished => println!("✅ Task completed normally"), EngineResult::Finished => println!("✅ Task completed normally"),
EngineResult::Paused => println!("⏸️ Task paused"), EngineResult::Paused => println!("⏸️ Task paused"),
crates/agen/examples/engine_cli.rs
+6 -4
View File
@@ -39,7 +39,7 @@ use tracing::info;
use tracing_subscriber::EnvFilter; use tracing_subscriber::EnvFilter;
use agen::{ use agen::{
Engine, Engine, History,
interceptor::{Interceptor, PostToolAction, ToolResultInfo}, interceptor::{Interceptor, PostToolAction, ToolResultInfo},
llm_client::{ llm_client::{
LlmClient, LlmClient,
@@ -474,9 +474,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
engine.set_interceptor(ToolResultPrinterPolicy::new(tool_call_names)); engine.set_interceptor(ToolResultPrinterPolicy::new(tool_call_names));
let mut history = History::new();
// One-shot mode // One-shot mode
if let Some(prompt) = args.prompt { if let Some(prompt) = args.prompt {
match engine.run(&prompt).await { match engine.run(&mut history, &prompt).await {
Ok(_) => {} Ok(_) => {}
Err(e) => { Err(e) => {
eprintln!("\n❌ Error: {}", e); eprintln!("\n❌ Error: {}", e);
@@ -500,7 +502,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
return Ok(()); return Ok(());
} }
let mut locked = match engine.run(first_input).await { let mut locked = match engine.run(&mut history, first_input).await {
Ok(out) => out.engine, Ok(out) => out.engine,
Err(e) => { Err(e) => {
eprintln!("\n❌ Error: {}", e); eprintln!("\n❌ Error: {}", e);
@@ -525,7 +527,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
break; break;
} }
match locked.run(input).await { match locked.run(&mut history, input).await {
Ok(_) => {} Ok(_) => {}
Err(e) => { Err(e) => {
eprintln!("\n❌ Error: {}", e); eprintln!("\n❌ Error: {}", e);
crates/agen/src/engine.rs
+180 -87
View File
@@ -7,7 +7,7 @@ use tokio::sync::mpsc;
use tracing::{debug, info, trace, warn}; use tracing::{debug, info, trace, warn};
use crate::{ use crate::{
Item, History, HistoryEntry, Item,
callback::{ callback::{
ClosureMetaHandler, ClosureTextBlockHandler, ClosureThinkingBlockHandler, ClosureMetaHandler, ClosureTextBlockHandler, ClosureThinkingBlockHandler,
ClosureToolUseBlockHandler, TextBlockScope, ThinkingBlockScope, ToolUseBlockScope, ClosureToolUseBlockHandler, TextBlockScope, ThinkingBlockScope, ToolUseBlockScope,
@@ -91,9 +91,9 @@ pub enum EngineResult {
/// Result of [`Engine::run`] or [`Engine::resume`]. /// Result of [`Engine::run`] or [`Engine::resume`].
/// ///
/// Contains the `Locked` Engine (ready for subsequent runs) and the outcome. /// Contains the `Locked` Engine (ready for subsequent runs) and the outcome.
pub struct EngineRunOutput<C: LlmClient> { pub struct EngineRunOutput<C: LlmClient, A = ()> {
/// The Engine, now in Locked state. /// The Engine, now in Locked state.
pub engine: Engine<C, Locked>, pub engine: Engine<C, Locked, A>,
/// Outcome of the turn. /// Outcome of the turn.
pub result: EngineResult, pub result: EngineResult,
} }
@@ -113,29 +113,31 @@ const MAX_STREAM_CONTINUATIONS: u32 = 3;
/// ///
/// # State Transitions (Type-state) /// # 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. /// - [`Locked`]: Cache-protected state. Prefix context is immutable; only `run()` / `resume()` are available.
/// ///
/// Calling `run()` on a `Mutable` Engine consumes it and returns a /// Calling `run()` on a `Mutable` Engine consumes it and returns a
/// `Locked` Engine together with the result. This ensures the /// `Locked` Engine together with the result. The engine borrows the caller's
/// cache prefix is fixed for optimal KV cache hit rate. /// [`History`](crate::History) only while running, so host annotations stay with
/// the host-owned history and are never projected to providers.
/// ///
/// ```ignore /// ```ignore
/// let mut history = History::new();
/// let mut engine = Engine::new(client) /// let mut engine = Engine::new(client)
/// .system_prompt("You are a helpful assistant."); /// .system_prompt("You are a helpful assistant.");
/// engine.register_tool(my_tool); /// engine.register_tool(my_tool);
/// ///
/// // Mutable::run() consumes self → EngineRunOutput { engine: Locked, result } /// // 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; /// let mut engine = out.engine;
/// ///
/// // Locked::run() borrows &mut self /// // Locked::run() borrows &mut self
/// engine.run("Follow-up").await?; /// engine.run(&mut history, "Follow-up").await?;
/// ///
/// // To edit between turns, unlock back to Mutable /// // To edit between turns, unlock back to Mutable
/// let mut engine = engine.unlock(); /// let mut engine = engine.unlock();
/// engine.truncate_history(5); /// history.truncate(5);
/// let out = engine.run("Continue").await?; /// let out = engine.run(&mut history, "Continue").await?;
/// let mut engine = out.engine; /// let mut engine = out.engine;
/// ``` /// ```
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
@@ -155,7 +157,7 @@ enum StreamCompletion {
Interrupted { reason: String }, Interrupted { reason: String },
} }
pub struct Engine<C: LlmClient, S: EngineState = Mutable> { pub struct Engine<C: LlmClient, S: EngineState = Mutable, A = ()> {
/// LLM client /// LLM client
client: C, client: C,
/// Retry policy for opening an LLM response stream. /// Retry policy for opening an LLM response stream.
@@ -175,8 +177,6 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
interceptor: Box<dyn Interceptor>, interceptor: Box<dyn Interceptor>,
/// System prompt /// System prompt
system_prompt: Option<String>, system_prompt: Option<String>,
/// Item history (owned by Engine)
history: Vec<Item>,
/// History length at lock time (only meaningful in Locked state) /// History length at lock time (only meaningful in Locked state)
locked_prefix_len: usize, locked_prefix_len: usize,
/// AgentTurn count across the lifetime of this Engine. /// AgentTurn count across the lifetime of this Engine.
@@ -266,10 +266,10 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
/// stable conversation identifier when the backend benefits from one. /// stable conversation identifier when the backend benefits from one.
cache_key: Option<String>, cache_key: Option<String>,
/// State marker /// State marker
_state: PhantomData<S>, _state: PhantomData<(S, A)>,
} }
impl<C: LlmClient, S: EngineState> Engine<C, S> { impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
fn reset_interruption_state(&mut self) { fn reset_interruption_state(&mut self) {
self.last_run_interrupted = false; self.last_run_interrupted = false;
} }
@@ -539,11 +539,15 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
fn append_history_items( fn append_history_items(
&mut self, &mut self,
history: &mut History<A>,
items: impl IntoIterator<Item = Item>, items: impl IntoIterator<Item = Item>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> Result<(), EngineError> { ) -> Result<(), EngineError> {
for item in items { for item in items {
self.emit_history_append(&item)?; self.emit_history_append(&item)?;
self.history.push(item); history
.append_with(item, annotate)
.map_err(EngineError::HistoryAppend)?;
} }
Ok(()) Ok(())
} }
@@ -650,9 +654,9 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
&self.client &self.client
} }
/// Get a reference to the history /// Borrow caller-owned annotated history entries.
pub fn history(&self) -> &[Item] { pub fn history<'h>(&self, history: &'h History<A>) -> &'h [HistoryEntry<A>] {
&self.history history.entries()
} }
/// Get a reference to the system prompt /// Get a reference to the system prompt
@@ -915,20 +919,20 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
} }
/// Check for pending tool calls (for resuming from Pause) /// Check for pending tool calls (for resuming from Pause)
fn get_pending_tool_calls(&self) -> Option<Vec<ToolCall>> { fn get_pending_tool_calls(&self, history: &History<A>) -> Option<Vec<ToolCall>> {
// Find the last ToolCall items that don't have corresponding ToolResult // Find the last ToolCall items that don't have corresponding ToolResult
let mut pending_calls = Vec::new(); let mut pending_calls = Vec::new();
let mut answered_call_ids = std::collections::HashSet::new(); let mut answered_call_ids = std::collections::HashSet::new();
// First pass: collect all answered call IDs // First pass: collect all answered call IDs
for item in &self.history { for item in history.items() {
if let Item::ToolResult { call_id, .. } = item { if let Item::ToolResult { call_id, .. } = item {
answered_call_ids.insert(call_id.clone()); answered_call_ids.insert(call_id.clone());
} }
} }
// Second pass: find unanswered tool calls // Second pass: find unanswered tool calls
for item in &self.history { for item in history.items() {
if let Item::ToolCall { if let Item::ToolCall {
call_id, call_id,
name, name,
@@ -1132,20 +1136,27 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
} }
/// Internal turn execution logic /// Internal turn execution logic
async fn run_turn_loop(&mut self) -> Result<EngineResult, EngineError> { async fn run_turn_loop(
&mut self,
history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> Result<EngineResult, EngineError> {
self.reset_interruption_state(); self.reset_interruption_state();
let tool_definitions = self.build_tool_definitions(); let tool_definitions = self.build_tool_definitions();
info!( info!(
item_count = self.history.len(), item_count = history.len(),
tool_count = tool_definitions.len(), tool_count = tool_definitions.len(),
"Starting engine run" "Starting engine run"
); );
// Resume pending tool calls from a previous Pause // 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"); 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); return Ok(result);
} }
} }
@@ -1192,13 +1203,13 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
.await .await
.map_err(EngineError::HistoryAppend)?; .map_err(EngineError::HistoryAppend)?;
if !pending.is_empty() { 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 // Clone the history into a per-request context. Everything
// below (prune projection, interceptor hooks) mutates only // below (prune projection, interceptor hooks) mutates only
// this clone, so the persistent `self.history` stays intact. // this clone, so the caller-owned `history` stays intact.
let mut request_context = self.history.clone(); let mut request_context = history.items_cloned();
// Prune projection: if both the config and the savings // Prune projection: if both the config and the savings
// estimator are configured, drop ToolResult.content from // estimator are configured, drop ToolResult.content from
@@ -1267,7 +1278,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
return Err(EngineError::Aborted(reason)); return Err(EngineError::Aborted(reason));
} }
PreRequestAction::YieldWith(items) => { PreRequestAction::YieldWith(items) => {
self.append_history_items(items.clone())?; self.append_history_items(history, items.clone(), annotate)?;
request_context.extend(items); request_context.extend(items);
info!("Yielded by interceptor after pre-request history append"); info!("Yielded by interceptor after pre-request history append");
for cb in &self.turn_end_cbs { for cb in &self.turn_end_cbs {
@@ -1285,7 +1296,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
return Ok(EngineResult::Yielded); return Ok(EngineResult::Yielded);
} }
PreRequestAction::ContinueWith(items) => { PreRequestAction::ContinueWith(items) => {
self.append_history_items(items.clone())?; self.append_history_items(history, items.clone(), annotate)?;
request_context.extend(items); request_context.extend(items);
} }
PreRequestAction::Continue => {} PreRequestAction::Continue => {}
@@ -1345,7 +1356,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
let assistant_items = let assistant_items =
self.build_assistant_items(&reasoning_items, &text_blocks, &[]); self.build_assistant_items(&reasoning_items, &text_blocks, &[]);
if !assistant_items.is_empty() { if !assistant_items.is_empty() {
self.append_history_items(assistant_items)?; self.append_history_items(history, assistant_items, annotate)?;
} }
self.emit_llm_continuation( self.emit_llm_continuation(
current_llm_call, current_llm_call,
@@ -1373,16 +1384,17 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
let tool_calls = self.tool_call_collector.take_collected(); let tool_calls = self.tool_call_collector.take_collected();
let assistant_items = let assistant_items =
self.build_assistant_items(&reasoning_items, &text_blocks, &tool_calls); 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() { 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 => { TurnEndAction::Finish => {
self.last_run_interrupted = false; self.last_run_interrupted = false;
return Ok(EngineResult::Finished); return Ok(EngineResult::Finished);
} }
TurnEndAction::ContinueWithMessages(additional) => { TurnEndAction::ContinueWithMessages(additional) => {
self.append_history_items(additional)?; self.append_history_items(history, additional, annotate)?;
continue; continue;
} }
TurnEndAction::Pause => { TurnEndAction::Pause => {
@@ -1392,7 +1404,10 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
} }
} }
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); return Ok(result);
} }
} }
@@ -1646,6 +1661,8 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
/// `None` if the turn loop should continue. /// `None` if the turn loop should continue.
async fn execute_and_commit_tools( async fn execute_and_commit_tools(
&mut self, &mut self,
history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
tool_calls: Vec<ToolCall>, tool_calls: Vec<ToolCall>,
) -> Result<Option<EngineResult>, EngineError> { ) -> Result<Option<EngineResult>, EngineError> {
match self.execute_tools(tool_calls).await { match self.execute_tools(tool_calls).await {
@@ -1665,7 +1682,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
result.attachments, result.attachments,
) )
}); });
self.append_history_items(items)?; self.append_history_items(history, items, annotate)?;
Ok(None) Ok(None)
} }
Err(err) => { Err(err) => {
@@ -1676,9 +1693,9 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
} }
} }
impl<C: LlmClient> Engine<C, Mutable> { impl<C: LlmClient, A> Engine<C, Mutable, A> {
/// Create a new Engine (in Mutable state) /// Create a new annotated Engine (in Mutable state).
pub fn new(client: C) -> Self { pub fn new_annotated(client: C) -> Self {
let text_block_collector = TextBlockCollector::new(); let text_block_collector = TextBlockCollector::new();
let tool_call_collector = ToolCallCollector::new(); let tool_call_collector = ToolCallCollector::new();
let thinking_block_collector = ThinkingBlockCollector::new(); let thinking_block_collector = ThinkingBlockCollector::new();
@@ -1700,7 +1717,6 @@ impl<C: LlmClient> Engine<C, Mutable> {
tool_server: ToolServer::new().handle(), tool_server: ToolServer::new().handle(),
interceptor: Box::new(DefaultInterceptor), interceptor: Box::new(DefaultInterceptor),
system_prompt: None, system_prompt: None,
history: Vec::new(),
locked_prefix_len: 0, locked_prefix_len: 0,
turn_count: 0, turn_count: 0,
active_run_turn_count: None, active_run_turn_count: None,
@@ -1861,36 +1877,38 @@ impl<C: LlmClient> Engine<C, Mutable> {
} }
} }
/// 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 /// This is not a history-growth API. Live append paths must use
/// [`append_history`](Self::append_history) so `on_history_append` observers /// [`append_history_with`](Self::append_history_with) so observers and the
/// see every inserted item. /// trusted annotation callback see every inserted item.
pub fn set_history(&mut self, items: Vec<Item>) { pub fn replace_history_entries(
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(
&mut self, &mut self,
history: &mut History<A>,
entries: Vec<HistoryEntry<A>>,
) -> Vec<HistoryEntry<A>> {
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<A>,
items: impl IntoIterator<Item = Item>, items: impl IntoIterator<Item = Item>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> Result<(), EngineError> { ) -> Result<(), EngineError> {
self.append_history_items(items) self.append_history_items(history, items, annotate)
} }
/// Truncate history without emitting append callbacks. /// Truncate caller-owned history without emitting append callbacks.
/// pub fn truncate_history(&mut self, history: &mut History<A>, len: usize) {
/// This is an edit operation, not a history-growth path. history.truncate(len);
pub fn truncate_history(&mut self, len: usize) {
self.history.truncate(len);
} }
/// Clear history /// Clear caller-owned history.
pub fn clear_history(&mut self) { pub fn clear_history(&mut self, history: &mut History<A>) {
self.history.clear(); history.clear();
} }
/// Set the turn count (for session restoration) /// Set the turn count (for session restoration)
@@ -1917,19 +1935,21 @@ impl<C: LlmClient> Engine<C, Mutable> {
self self
} }
/// Execute a turn, consuming self and transitioning to Locked. /// Run the engine with one user input, appending to caller-owned history.
/// ///
/// This is the primary entry point for first use. Equivalent to /// The trusted `annotate` callback is invoked after append observers and before
/// `self.lock()` followed by `locked.run(user_input)`. /// each new item becomes live in `history`. Providers, token counters, pruners,
/// /// and interceptors receive only the `Item` projection.
/// Subsequent runs can call [`Engine::run`] directly. pub async fn run_with_annotation(
/// To edit state between turns, call [`unlock()`](Engine::unlock) first.
pub async fn run(
self, self,
history: &mut History<A>,
user_input: impl Into<String>, user_input: impl Into<String>,
) -> Result<EngineRunOutput<C>, EngineError> { annotate: &mut impl FnMut(&Item) -> Result<A, String>,
let mut locked = self.lock(); ) -> Result<EngineRunOutput<C, A>, EngineError> {
let result = locked.run(user_input).await?; let mut locked = self.lock(history);
let result = locked
.run_with_annotation(history, user_input, annotate)
.await?;
Ok(EngineRunOutput { Ok(EngineRunOutput {
engine: locked, engine: locked,
result, result,
@@ -1939,9 +1959,13 @@ impl<C: LlmClient> Engine<C, Mutable> {
/// Resume from Paused, consuming self and transitioning to Locked. /// Resume from Paused, consuming self and transitioning to Locked.
/// ///
/// Used after `unlock()` → edit → resume. /// Used after `unlock()` → edit → resume.
pub async fn resume(self) -> Result<EngineRunOutput<C>, EngineError> { pub async fn resume_with_annotation(
let mut locked = self.lock(); self,
let result = locked.resume().await?; history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> Result<EngineRunOutput<C, A>, EngineError> {
let mut locked = self.lock(history);
let result = locked.resume_with_annotation(history, annotate).await?;
Ok(EngineRunOutput { Ok(EngineRunOutput {
engine: locked, engine: locked,
result, result,
@@ -1961,9 +1985,9 @@ impl<C: LlmClient> Engine<C, Mutable> {
/// # Panics /// # Panics
/// ///
/// Panics if a pending tool factory produces a duplicate name. /// Panics if a pending tool factory produces a duplicate name.
pub fn lock(self) -> Engine<C, Locked> { pub fn lock(self, history: &History<A>) -> Engine<C, Locked, A> {
self.tool_server.flush_pending(); self.tool_server.flush_pending();
let locked_prefix_len = self.history.len(); let locked_prefix_len = history.len();
Engine { Engine {
client: self.client, client: self.client,
retry_policy: self.retry_policy, retry_policy: self.retry_policy,
@@ -1974,7 +1998,6 @@ impl<C: LlmClient> Engine<C, Mutable> {
tool_server: self.tool_server, tool_server: self.tool_server,
interceptor: self.interceptor, interceptor: self.interceptor,
system_prompt: self.system_prompt, system_prompt: self.system_prompt,
history: self.history,
locked_prefix_len, locked_prefix_len,
turn_count: self.turn_count, turn_count: self.turn_count,
active_run_turn_count: self.active_run_turn_count, active_run_turn_count: self.active_run_turn_count,
@@ -2009,14 +2032,62 @@ impl<C: LlmClient> Engine<C, Mutable> {
} }
} }
impl<C: LlmClient> Engine<C, Locked> { fn unit_history_annotation(_: &Item) -> Result<(), String> {
Ok(())
}
impl<C: LlmClient> Engine<C, Mutable, ()> {
/// 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<Item = Item>,
) -> 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<Item>) {
history.replace_items(items);
}
/// Run using unit annotations.
pub async fn run(
self,
history: &mut History<()>,
user_input: impl Into<String>,
) -> Result<EngineRunOutput<C>, EngineError> {
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<()>,
) -> Result<EngineRunOutput<C>, EngineError> {
let mut annotate = unit_history_annotation;
self.resume_with_annotation(history, &mut annotate).await
}
}
impl<C: LlmClient, A> Engine<C, Locked, A> {
/// Execute a turn /// Execute a turn
/// ///
/// Adds a new user message to history and sends a request to the LLM. /// Adds a new user message to history and sends a request to the LLM.
/// Automatically loops if there are tool calls. /// Automatically loops if there are tool calls.
pub async fn run( pub async fn run_with_annotation(
&mut self, &mut self,
history: &mut History<A>,
user_input: impl Into<String>, user_input: impl Into<String>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> Result<EngineResult, EngineError> { ) -> Result<EngineResult, EngineError> {
// Supplying new user input abandons any paused/yielded logical run. // Supplying new user input abandons any paused/yielded logical run.
self.active_run_turn_count = None; self.active_run_turn_count = None;
@@ -2033,12 +2104,12 @@ impl<C: LlmClient> Engine<C, Locked> {
PromptAction::Continue => Vec::new(), PromptAction::Continue => Vec::new(),
PromptAction::ContinueWith(items) => items, 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() { if !extras.is_empty() {
self.append_history_items(extras)?; self.append_history_items(history, extras, annotate)?;
} }
self.start_logical_run(); 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; let result = self.finalize_interruption(result).await;
self.finish_logical_run(&result); self.finish_logical_run(&result);
result result
@@ -2047,10 +2118,14 @@ impl<C: LlmClient> Engine<C, Locked> {
/// Resume execution (from Paused state) /// Resume execution (from Paused state)
/// ///
/// Resumes turn processing from current state without adding a new user message. /// Resumes turn processing from current state without adding a new user message.
pub async fn resume(&mut self) -> Result<EngineResult, EngineError> { pub async fn resume_with_annotation(
&mut self,
history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> Result<EngineResult, EngineError> {
self.reset_interruption_state(); self.reset_interruption_state();
self.ensure_logical_run(); 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; let result = self.finalize_interruption(result).await;
self.finish_logical_run(&result); self.finish_logical_run(&result);
result result
@@ -2065,7 +2140,7 @@ impl<C: LlmClient> Engine<C, Locked> {
/// ///
/// Note: After this operation, subsequent requests may not hit the cache. /// Note: After this operation, subsequent requests may not hit the cache.
/// Use only when you need to edit history. /// Use only when you need to edit history.
pub fn unlock(self) -> Engine<C, Mutable> { pub fn unlock(self) -> Engine<C, Mutable, A> {
Engine { Engine {
client: self.client, client: self.client,
retry_policy: self.retry_policy, retry_policy: self.retry_policy,
@@ -2076,7 +2151,6 @@ impl<C: LlmClient> Engine<C, Locked> {
tool_server: self.tool_server, tool_server: self.tool_server,
interceptor: self.interceptor, interceptor: self.interceptor,
system_prompt: self.system_prompt, system_prompt: self.system_prompt,
history: self.history,
locked_prefix_len: 0, locked_prefix_len: 0,
turn_count: self.turn_count, turn_count: self.turn_count,
active_run_turn_count: self.active_run_turn_count, active_run_turn_count: self.active_run_turn_count,
@@ -2111,6 +2185,25 @@ impl<C: LlmClient> Engine<C, Locked> {
} }
} }
impl<C: LlmClient> Engine<C, Locked, ()> {
/// Run another turn using unit annotations.
pub async fn run(
&mut self,
history: &mut History<()>,
user_input: impl Into<String>,
) -> Result<EngineResult, EngineError> {
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<()>) -> Result<EngineResult, EngineError> {
let mut annotate = unit_history_annotation;
self.resume_with_annotation(history, &mut annotate).await
}
}
enum FirstStreamEvent { enum FirstStreamEvent {
Ready(ResponseStream), Ready(ResponseStream),
Empty(ResponseStream), Empty(ResponseStream),
crates/agen/src/history.rs
+199
View File
@@ -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<A = ()> {
/// Provider/model-visible conversation item.
pub item: Item,
/// Host-domain metadata kept with the item and never projected to providers.
pub annotation: A,
}
impl<A> HistoryEntry<A> {
/// 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<A = ()> {
entries: Vec<HistoryEntry<A>>,
}
impl<A> History<A> {
/// 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<HistoryEntry<A>>) -> Self {
Self { entries }
}
/// Replace all entries as one restore/rebuild operation and return the old entries.
pub fn replace_entries(&mut self, entries: Vec<HistoryEntry<A>>) -> Vec<HistoryEntry<A>> {
std::mem::replace(&mut self.entries, entries)
}
/// Borrow annotated entries.
pub fn entries(&self) -> &[HistoryEntry<A>] {
&self.entries
}
/// Mutably borrow annotated entries for host-owned rebuild operations.
pub fn entries_mut(&mut self) -> &mut [HistoryEntry<A>] {
&mut self.entries
}
/// Consume the history into annotated entries.
pub fn into_entries(self) -> Vec<HistoryEntry<A>> {
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<Item = &HistoryEntry<A>> {
self.entries.iter()
}
/// Iterate over provider-visible items only.
pub fn items(&self) -> impl ExactSizeIterator<Item = &Item> {
self.entries.iter().map(|entry| &entry.item)
}
/// Clone provider-visible items into a request-local projection.
pub fn items_cloned(&self) -> Vec<Item> {
self.items().cloned().collect()
}
/// Append an already annotated entry.
pub fn push_entry(&mut self, entry: HistoryEntry<A>) {
self.entries.push(entry);
}
/// Append many already annotated entries.
pub fn extend_entries(&mut self, entries: impl IntoIterator<Item = HistoryEntry<A>>) {
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<A, String>,
) -> 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<Item = Item>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> 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<Item>) -> 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<Item>) -> Vec<HistoryEntry<()>> {
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<Item = Item>) {
self.entries
.extend(items.into_iter().map(HistoryEntry::from_item));
}
}
impl<A> IntoIterator for History<A> {
type Item = HistoryEntry<A>;
type IntoIter = std::vec::IntoIter<HistoryEntry<A>>;
fn into_iter(self) -> Self::IntoIter {
self.entries.into_iter()
}
}
impl<'a, A> IntoIterator for &'a History<A> {
type Item = &'a HistoryEntry<A>;
type IntoIter = std::slice::Iter<'a, HistoryEntry<A>>;
fn into_iter(self) -> Self::IntoIter {
self.entries.iter()
}
}
crates/agen/src/lib.rs
+2
View File
@@ -2,6 +2,7 @@
mod engine; mod engine;
mod handler; mod handler;
mod history;
mod message; mod message;
pub(crate) mod callback; pub(crate) mod callback;
@@ -24,6 +25,7 @@ pub use engine::{
ToolRegistryError, ToolRegistryError,
}; };
pub use handler::ToolUseBlockStart; pub use handler::ToolUseBlockStart;
pub use history::{History, HistoryEntry};
pub use interceptor::Interceptor; pub use interceptor::Interceptor;
pub use message::{ContentPart, Item, Message, Role}; pub use message::{ContentPart, Item, Message, Role};
pub use tool::{ToolCall, ToolExecutionContext, ToolOutputLimits, ToolResult}; pub use tool::{ToolCall, ToolExecutionContext, ToolOutputLimits, ToolResult};
crates/agen/src/state.rs
+1 -1
View File
@@ -19,7 +19,7 @@ mod private {
/// - Editing message history (add, delete, clear) /// - Editing message history (add, delete, clear)
/// - Registering tools and hooks /// - Registering tools and hooks
/// ///
/// Can transition to [`Locked`] state via `Engine::lock()`. /// Can transition to [`Locked`] state via `Engine::lock(&history)`.
/// ///
/// # Examples /// # Examples
/// ///
crates/agen/tests/annotated_history_test.rs
+85
View File
@@ -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<Event> {
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::<String>::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::<usize>::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");
}
crates/agen/tests/callback_test.rs
+15 -8
View File
@@ -8,11 +8,11 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::Duration; use std::time::Duration;
use agen::Engine;
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent as ClientStatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent as ClientStatusEvent};
use agen::llm_client::retry::RetryPolicy; use agen::llm_client::retry::RetryPolicy;
use agen::llm_client::{ClientError, LlmClient, Request, ResponseStream}; use agen::llm_client::{ClientError, LlmClient, Request, ResponseStream};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, History};
use async_trait::async_trait; use async_trait::async_trait;
use common::MockLlmClient; use common::MockLlmClient;
@@ -58,6 +58,7 @@ async fn test_callback_llm_retry_event() {
max_attempts: 2, max_attempts: 2,
total_timeout: Duration::from_secs(1), total_timeout: Duration::from_secs(1),
}); });
let mut history: History = History::new();
let notices = Arc::new(Mutex::new(Vec::new())); let notices = Arc::new(Mutex::new(Vec::new()));
let sink = notices.clone(); let sink = notices.clone();
@@ -65,7 +66,7 @@ async fn test_callback_llm_retry_event() {
sink.lock().unwrap().push((llm_call, notice.clone())); sink.lock().unwrap().push((llm_call, notice.clone()));
}); });
let result = engine.run("retry once").await; let result = engine.run(&mut history, "retry once").await;
assert!(result.is_ok(), "engine should succeed after one retry"); assert!(result.is_ok(), "engine should succeed after one retry");
let notices = notices.lock().unwrap(); let notices = notices.lock().unwrap();
@@ -91,6 +92,7 @@ async fn test_callback_text_block_events() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let text_deltas = Arc::new(Mutex::new(Vec::new())); let text_deltas = Arc::new(Mutex::new(Vec::new()));
let text_completes = Arc::new(Mutex::new(Vec::new())); let text_completes = Arc::new(Mutex::new(Vec::new()));
@@ -109,7 +111,7 @@ async fn test_callback_text_block_events() {
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineResult)
let result = engine.run("Greet me").await; let result = engine.run(&mut history, "Greet me").await;
assert!(result.is_ok(), "Engine should complete"); assert!(result.is_ok(), "Engine should complete");
let deltas = text_deltas.lock().unwrap(); let deltas = text_deltas.lock().unwrap();
@@ -137,6 +139,7 @@ async fn test_callback_tool_call_complete() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); 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_starts = Arc::new(Mutex::new(Vec::<(String, String)>::new()));
let tool_completes = Arc::new(Mutex::new(Vec::new())); let tool_completes = Arc::new(Mutex::new(Vec::new()));
@@ -155,7 +158,7 @@ async fn test_callback_tool_call_complete() {
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineResult)
let _ = engine.run("Weather please").await; let _ = engine.run(&mut history, "Weather please").await;
let starts = tool_starts.lock().unwrap(); let starts = tool_starts.lock().unwrap();
assert_eq!(starts.len(), 1); assert_eq!(starts.len(), 1);
@@ -183,6 +186,7 @@ async fn test_callback_turn_events() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let turn_starts = Arc::new(Mutex::new(Vec::new())); let turn_starts = Arc::new(Mutex::new(Vec::new()));
let turn_ends = Arc::new(Mutex::new(Vec::new())); let turn_ends = Arc::new(Mutex::new(Vec::new()));
@@ -198,7 +202,7 @@ async fn test_callback_turn_events() {
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineResult)
let result = engine.run("Do something").await; let result = engine.run(&mut history, "Do something").await;
assert!(result.is_ok()); assert!(result.is_ok());
let starts = turn_starts.lock().unwrap(); let starts = turn_starts.lock().unwrap();
@@ -254,6 +258,7 @@ async fn test_callback_tool_result_events() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine.register_tool(fixed_tool( engine.register_tool(fixed_tool(
"fixed", "fixed",
@@ -276,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(); let observed = captured.lock().unwrap();
assert_eq!(observed.len(), 1); assert_eq!(observed.len(), 1);
@@ -330,6 +335,7 @@ async fn test_callback_tool_result_error_path() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine.register_tool(erroring_tool("erroring", "boom")); engine.register_tool(erroring_tool("erroring", "boom"));
@@ -345,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(); let observed = captured.lock().unwrap();
assert_eq!(observed.len(), 1); assert_eq!(observed.len(), 1);
@@ -374,6 +380,7 @@ async fn test_callback_usage_events() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let usage_events = Arc::new(Mutex::new(Vec::new())); let usage_events = Arc::new(Mutex::new(Vec::new()));
@@ -383,7 +390,7 @@ async fn test_callback_usage_events() {
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineResult)
let _ = engine.run("Hello").await; let _ = engine.run(&mut history, "Hello").await;
let usages = usage_events.lock().unwrap(); let usages = usage_events.lock().unwrap();
assert_eq!(usages.len(), 1); assert_eq!(usages.len(), 1);
crates/agen/tests/engine_fixtures.rs
+9 -4
View File
@@ -9,8 +9,8 @@ use std::path::Path;
use std::sync::Arc; use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::atomic::{AtomicUsize, Ordering};
use agen::Engine;
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, History};
use async_trait::async_trait; use async_trait::async_trait;
use common::MockLlmClient; use common::MockLlmClient;
@@ -134,9 +134,10 @@ async fn test_engine_simple_text_response() {
let client = MockLlmClient::from_fixture(&fixture_path).unwrap(); let client = MockLlmClient::from_fixture(&fixture_path).unwrap();
let engine = Engine::new(client); let engine = Engine::new(client);
let mut history: History = History::new();
// Send a simple message (Mutable::run consumes self, returns tuple) // Send a simple message (Mutable::run consumes self, returns tuple)
let result = engine.run("Hello").await; let result = engine.run(&mut history, "Hello").await;
assert!(result.is_ok(), "Engine should complete successfully"); assert!(result.is_ok(), "Engine should complete successfully");
} }
@@ -156,6 +157,7 @@ async fn test_engine_tool_call() {
let client = MockLlmClient::from_fixture(&fixture_path).unwrap(); let client = MockLlmClient::from_fixture(&fixture_path).unwrap();
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
// Register tool // Register tool
let weather_tool = MockWeatherTool::new(); let weather_tool = MockWeatherTool::new();
@@ -163,7 +165,9 @@ async fn test_engine_tool_call() {
engine.register_tool(weather_tool.definition()); engine.register_tool(weather_tool.definition());
// Send message (Mutable::run consumes self, returns tuple) // 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 // Verify tool was called
// Note: max_turns=1 so no request is sent after tool result // Note: max_turns=1 so no request is sent after tool result
@@ -195,9 +199,10 @@ async fn test_engine_with_programmatic_events() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let engine = Engine::new(client); let engine = Engine::new(client);
let mut history: History = History::new();
// Mutable::run consumes self, returns tuple // Mutable::run consumes self, returns tuple
let result = engine.run("Greet me").await; let result = engine.run(&mut history, "Greet me").await;
assert!(result.is_ok(), "Engine should complete successfully"); assert!(result.is_ok(), "Engine should complete successfully");
} }
crates/agen/tests/engine_state_test.rs
+148 -87
View File
@@ -14,7 +14,7 @@ use agen::interceptor::{
}; };
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, EngineError, EngineResult}; use agen::{Engine, EngineError, EngineResult, History};
use async_trait::async_trait; use async_trait::async_trait;
use common::MockLlmClient; use common::MockLlmClient;
@@ -42,36 +42,37 @@ fn test_mutable_set_system_prompt() {
fn test_mutable_history_manipulation() { fn test_mutable_history_manipulation() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
// Initial state is empty // Initial state is empty
assert!(engine.history().is_empty()); assert!(history.is_empty());
// Add to history // Add to history
engine engine
.append_history(vec![Item::user_message("Hello")]) .append_history(&mut history, vec![Item::user_message("Hello")])
.unwrap(); .unwrap();
engine engine
.append_history(vec![Item::assistant_message("Hi there!")]) .append_history(&mut history, vec![Item::assistant_message("Hi there!")])
.unwrap(); .unwrap();
assert_eq!(engine.history().len(), 2); assert_eq!(history.len(), 2);
// Append to history via the callback-aware API. // Append to history via the callback-aware API.
engine engine
.append_history(vec![Item::user_message("How are you?")]) .append_history(&mut history, vec![Item::user_message("How are you?")])
.unwrap(); .unwrap();
assert_eq!(engine.history().len(), 3); assert_eq!(history.len(), 3);
// Clear history // Clear history
engine.clear_history(); engine.clear_history(&mut history);
assert!(engine.history().is_empty()); assert!(history.is_empty());
// Set history // Set history
let items = vec![ let items = vec![
Item::user_message("Test"), Item::user_message("Test"),
Item::assistant_message("Response"), Item::assistant_message("Response"),
]; ];
engine.set_history(items); engine.set_history(&mut history, items);
assert_eq!(engine.history().len(), 2); assert_eq!(history.len(), 2);
} }
/// Verify that Engine can be constructed using builder pattern /// Verify that Engine can be constructed using builder pattern
@@ -79,9 +80,10 @@ fn test_mutable_history_manipulation() {
fn test_mutable_builder_pattern() { fn test_mutable_builder_pattern() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let engine = Engine::new(client).system_prompt("System prompt"); let engine = Engine::new(client).system_prompt("System prompt");
let history: History = History::new();
assert_eq!(engine.get_system_prompt(), Some("System prompt")); 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. /// Verify that multiple items can be added with append_history and callbacks fire.
@@ -91,6 +93,7 @@ fn test_mutable_append_history() {
let observed = Arc::new(Mutex::new(Vec::new())); let observed = Arc::new(Mutex::new(Vec::new()));
let observed_for_callback = Arc::clone(&observed); let observed_for_callback = Arc::clone(&observed);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine.on_history_append(move |item| { engine.on_history_append(move |item| {
if let Some(text) = item.as_text() { if let Some(text) = item.as_text() {
observed_for_callback.lock().unwrap().push(text.to_string()); observed_for_callback.lock().unwrap().push(text.to_string());
@@ -99,18 +102,21 @@ fn test_mutable_append_history() {
}); });
engine engine
.append_history(vec![Item::user_message("First")]) .append_history(&mut history, vec![Item::user_message("First")])
.unwrap(); .unwrap();
engine engine
.append_history(vec![ .append_history(
Item::assistant_message("Response 1"), &mut history,
Item::user_message("Second"), vec![
Item::assistant_message("Response 2"), Item::assistant_message("Response 1"),
]) Item::user_message("Second"),
Item::assistant_message("Response 2"),
],
)
.unwrap(); .unwrap();
assert_eq!(engine.history().len(), 4); assert_eq!(history.len(), 4);
assert_eq!( assert_eq!(
observed.lock().unwrap().as_slice(), observed.lock().unwrap().as_slice(),
["First", "Response 1", "Second", "Response 2"] ["First", "Response 1", "Second", "Response 2"]
@@ -185,6 +191,7 @@ async fn history_append_failure_stops_before_tool_execution() {
]); ]);
let tool = CountingTool::new("count_tool"); let tool = CountingTool::new("count_tool");
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine.register_tool(tool.definition()); engine.register_tool(tool.definition());
engine.on_history_append(|item| { engine.on_history_append(|item| {
if item.is_tool_call() { if item.is_tool_call() {
@@ -194,15 +201,15 @@ async fn history_append_failure_stops_before_tool_execution() {
} }
}); });
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
let error = engine.run("use the tool").await.unwrap_err(); let error = engine.run(&mut history, "use the tool").await.unwrap_err();
assert!( assert!(
matches!(error, EngineError::HistoryAppend(ref message) if message == "simulated ENOSPC") matches!(error, EngineError::HistoryAppend(ref message) if message == "simulated ENOSPC")
); );
assert_eq!(tool.call_count(), 0); assert_eq!(tool.call_count(), 0);
assert_eq!(engine.history().len(), 1); assert_eq!(history.len(), 1);
assert_eq!(engine.history()[0].as_text(), Some("use the tool")); assert_eq!(history.entries()[0].item.as_text(), Some("use the tool"));
} }
// ============================================================================= // =============================================================================
@@ -214,21 +221,22 @@ async fn history_append_failure_stops_before_tool_execution() {
fn test_lock_transition() { fn test_lock_transition() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine.set_system_prompt("System"); engine.set_system_prompt("System");
engine engine
.append_history(vec![Item::user_message("Hello")]) .append_history(&mut history, vec![Item::user_message("Hello")])
.unwrap(); .unwrap();
engine engine
.append_history(vec![Item::assistant_message("Hi")]) .append_history(&mut history, vec![Item::assistant_message("Hi")])
.unwrap(); .unwrap();
// Lock // Lock
let locked_engine = engine.lock(); let locked_engine = engine.lock(&history);
// History and system prompt are still accessible in Locked state // History and system prompt are still accessible in Locked state
assert_eq!(locked_engine.get_system_prompt(), Some("System")); 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); assert_eq!(locked_engine.locked_prefix_len(), 2);
} }
@@ -237,21 +245,22 @@ fn test_lock_transition() {
fn test_unlock_transition() { fn test_unlock_transition() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine engine
.append_history(vec![Item::user_message("Hello")]) .append_history(&mut history, vec![Item::user_message("Hello")])
.unwrap(); .unwrap();
let locked_engine = engine.lock(); let locked_engine = engine.lock(&history);
// Unlock // Unlock
let mut engine = locked_engine.unlock(); let mut engine = locked_engine.unlock();
// History operations are available again in Mutable state // History operations are available again in Mutable state
engine engine
.append_history(vec![Item::assistant_message("Hi")]) .append_history(&mut history, vec![Item::assistant_message("Hi")])
.unwrap(); .unwrap();
engine.clear_history(); engine.clear_history(&mut history);
assert!(engine.history().is_empty()); assert!(history.is_empty());
} }
// ============================================================================= // =============================================================================
@@ -272,20 +281,20 @@ async fn test_mutable_run_updates_history() -> Result<(), EngineError> {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let engine = Engine::new(client); let engine = Engine::new(client);
let mut history: History = History::new();
// Execute (Mutable::run consumes self, returns EngineRunOutput) // Execute (Mutable::run consumes self, returns EngineRunOutput)
let out = engine.run("Hi there").await?; let _out = engine.run(&mut history, "Hi there").await?;
let engine = out.engine;
// History is updated // History is updated
let history = engine.history(); let entries = history.entries();
assert_eq!(history.len(), 2); // user + assistant assert_eq!(history.len(), 2); // user + assistant
// User message // User message
assert_eq!(history[0].as_text(), Some("Hi there")); assert_eq!(entries[0].item.as_text(), Some("Hi there"));
// Assistant message // 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(()) Ok(())
} }
@@ -316,35 +325,36 @@ async fn test_locked_multi_turn_history_accumulation() {
]); ]);
let engine = Engine::new(client).system_prompt("You are helpful."); let engine = Engine::new(client).system_prompt("You are helpful.");
let mut history: History = History::new();
// Lock (after setting system prompt) // 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 assert_eq!(locked_engine.locked_prefix_len(), 0); // No items yet
// Turn 1 // Turn 1
let result1 = locked_engine.run("Hello!").await; let result1 = locked_engine.run(&mut history, "Hello!").await;
assert!(result1.is_ok()); assert!(result1.is_ok());
assert_eq!(locked_engine.history().len(), 2); // user + assistant assert_eq!(history.len(), 2); // user + assistant
// Turn 2 // Turn 2
let result2 = locked_engine.run("Can you help me?").await; let result2 = locked_engine.run(&mut history, "Can you help me?").await;
assert!(result2.is_ok()); assert!(result2.is_ok());
assert_eq!(locked_engine.history().len(), 4); // 2 * (user + assistant) assert_eq!(history.len(), 4); // 2 * (user + assistant)
// Verify history contents // Verify history contents
let history = locked_engine.history(); let entries = history.entries();
// Turn 1 user message // 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 // 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 // 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 // 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 /// Verify that locked_prefix_len correctly records history length at lock time
@@ -370,26 +380,36 @@ async fn test_locked_prefix_len_tracking() {
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
// Add items beforehand // Add items beforehand
engine engine
.append_history(vec![Item::user_message("Pre-existing message 1")]) .append_history(
&mut history,
vec![Item::user_message("Pre-existing message 1")],
)
.unwrap(); .unwrap();
engine engine
.append_history(vec![Item::assistant_message("Pre-existing response 1")]) .append_history(
&mut history,
vec![Item::assistant_message("Pre-existing response 1")],
)
.unwrap(); .unwrap();
assert_eq!(engine.history().len(), 2); assert_eq!(history.len(), 2);
// Lock // 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 assert_eq!(locked_engine.locked_prefix_len(), 2); // 2 items at lock time
// Execute turn // Execute turn
locked_engine.run("New message").await.unwrap(); locked_engine
.run(&mut history, "New message")
.await
.unwrap();
// History grows but locked_prefix_len remains unchanged // 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 assert_eq!(locked_engine.locked_prefix_len(), 2); // Unchanged
} }
@@ -416,18 +436,19 @@ async fn test_turn_count_increment() -> Result<(), EngineError> {
]); ]);
let engine = Engine::new(client); let engine = Engine::new(client);
let mut history: History = History::new();
assert_eq!(engine.turn_count(), 0); assert_eq!(engine.turn_count(), 0);
assert_eq!(engine.llm_call_count(), 0); assert_eq!(engine.llm_call_count(), 0);
// First run consumes Mutable, returns EngineRunOutput // 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); assert_eq!(engine.turn_count(), 1);
// Retry not yet implemented → AgentTurn:LlmCall is 1:1. // Retry not yet implemented → AgentTurn:LlmCall is 1:1.
assert_eq!(engine.llm_call_count(), 1); assert_eq!(engine.llm_call_count(), 1);
// Subsequent runs on Locked take &mut self // 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.turn_count(), 2);
assert_eq!(engine.llm_call_count(), 2); assert_eq!(engine.llm_call_count(), 2);
@@ -447,28 +468,29 @@ async fn test_unlock_edit_relock() {
]]); ]]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine engine
.append_history(vec![ .append_history(
Item::user_message("Hello"), &mut history,
Item::assistant_message("Hi"), vec![Item::user_message("Hello"), Item::assistant_message("Hi")],
]) )
.unwrap(); .unwrap();
// Lock -> Unlock // Lock -> Unlock
let locked = engine.lock(); let locked = engine.lock(&history);
assert_eq!(locked.locked_prefix_len(), 2); assert_eq!(locked.locked_prefix_len(), 2);
let mut unlocked = locked.unlock(); let mut unlocked = locked.unlock();
// Edit history // Edit history
unlocked.clear_history(); unlocked.clear_history(&mut history);
unlocked unlocked
.append_history(vec![Item::user_message("Fresh start")]) .append_history(&mut history, vec![Item::user_message("Fresh start")])
.unwrap(); .unwrap();
// Re-lock // Re-lock
let relocked = unlocked.lock(); let relocked = unlocked.lock(&history);
assert_eq!(relocked.history().len(), 1); assert_eq!(history.len(), 1);
assert_eq!(relocked.locked_prefix_len(), 1); assert_eq!(relocked.locked_prefix_len(), 1);
} }
@@ -511,19 +533,23 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let tool_a = CountingTool::new("tool_a"); let tool_a = CountingTool::new("tool_a");
engine.register_tool(tool_a.definition()); engine.register_tool(tool_a.definition());
let mut locked = engine.lock(); let mut locked = engine.lock(&history);
locked.run("first").await.expect("first run"); locked.run(&mut history, "first").await.expect("first run");
assert_eq!(tool_a.call_count(), 1, "tool_a should be called once"); assert_eq!(tool_a.call_count(), 1, "tool_a should be called once");
let mut unlocked = locked.unlock(); let mut unlocked = locked.unlock();
let tool_b = CountingTool::new("tool_b"); let tool_b = CountingTool::new("tool_b");
unlocked.register_tool(tool_b.definition()); unlocked.register_tool(tool_b.definition());
let mut relocked = unlocked.lock(); let mut relocked = unlocked.lock(&history);
relocked.run("second").await.expect("second run"); 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_a.call_count(), 1, "tool_a should not be called again");
assert_eq!(tool_b.call_count(), 1, "tool_b should be called once"); assert_eq!(tool_b.call_count(), 1, "tool_b should be called once");
@@ -538,8 +564,9 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
fn test_system_prompt_preserved_in_locked_state() { fn test_system_prompt_preserved_in_locked_state() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let engine = Engine::new(client).system_prompt("Important system prompt"); 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")); assert_eq!(locked.get_system_prompt(), Some("Important system prompt"));
let unlocked = locked.unlock(); let unlocked = locked.unlock();
@@ -554,14 +581,15 @@ fn test_system_prompt_preserved_in_locked_state() {
fn test_system_prompt_change_after_unlock() { fn test_system_prompt_change_after_unlock() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let engine = Engine::new(client).system_prompt("Original prompt"); 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(); let mut unlocked = locked.unlock();
unlocked.set_system_prompt("New prompt"); unlocked.set_system_prompt("New prompt");
assert_eq!(unlocked.get_system_prompt(), Some("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")); assert_eq!(relocked.get_system_prompt(), Some("New prompt"));
} }
@@ -625,14 +653,21 @@ impl Interceptor for ContinueTurnOnce {
async fn max_turns_is_scoped_to_each_fresh_run() { async fn max_turns_is_scoped_to_each_fresh_run() {
let responses = vec![completed_text_events(), completed_text_events()]; let responses = vec![completed_text_events(), completed_text_events()];
let mut engine = Engine::new(MockLlmClient::with_responses(responses)); let mut engine = Engine::new(MockLlmClient::with_responses(responses));
let mut history: History = History::new();
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!(engine.run("first").await.unwrap(), EngineResult::Finished); assert_eq!(
engine.run(&mut history, "first").await.unwrap(),
EngineResult::Finished
);
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(engine.run("second").await.unwrap(), EngineResult::Finished); assert_eq!(
engine.run(&mut history, "second").await.unwrap(),
EngineResult::Finished
);
assert_eq!(engine.turn_count(), 2); assert_eq!(engine.turn_count(), 2);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
} }
@@ -640,17 +675,24 @@ async fn max_turns_is_scoped_to_each_fresh_run() {
#[tokio::test] #[tokio::test]
async fn yielded_resume_keeps_the_same_unspent_turn_budget() { async fn yielded_resume_keeps_the_same_unspent_turn_budget() {
let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
let mut history: History = History::new();
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
engine.set_interceptor(YieldOnce { engine.set_interceptor(YieldOnce {
calls: AtomicUsize::new(0), calls: AtomicUsize::new(0),
}); });
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!(engine.run("start").await.unwrap(), EngineResult::Yielded); assert_eq!(
engine.run(&mut history, "start").await.unwrap(),
EngineResult::Yielded
);
assert_eq!(engine.turn_count(), 0); assert_eq!(engine.turn_count(), 0);
assert_eq!(engine.active_run_turn_count(), Some(0)); assert_eq!(engine.active_run_turn_count(), Some(0));
assert_eq!(engine.resume().await.unwrap(), EngineResult::Finished); assert_eq!(
engine.resume(&mut history).await.unwrap(),
EngineResult::Finished
);
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
} }
@@ -667,19 +709,26 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
]; ];
let tool = CountingTool::new("count_tool"); let tool = CountingTool::new("count_tool");
let mut engine = Engine::new(MockLlmClient::new(events)); let mut engine = Engine::new(MockLlmClient::new(events));
let mut history: History = History::new();
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
engine.register_tool(tool.definition()); engine.register_tool(tool.definition());
engine.set_interceptor(PauseToolOnce { engine.set_interceptor(PauseToolOnce {
calls: AtomicUsize::new(0), calls: AtomicUsize::new(0),
}); });
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!(engine.run("call it").await.unwrap(), EngineResult::Paused); assert_eq!(
engine.run(&mut history, "call it").await.unwrap(),
EngineResult::Paused
);
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), Some(1)); assert_eq!(engine.active_run_turn_count(), Some(1));
assert_eq!(tool.call_count(), 0); assert_eq!(tool.call_count(), 0);
assert_eq!(engine.resume().await.unwrap(), EngineResult::LimitReached); assert_eq!(
engine.resume(&mut history).await.unwrap(),
EngineResult::LimitReached
);
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(tool.call_count(), 1, "the consumed turn's tool still runs"); assert_eq!(tool.call_count(), 1, "the consumed turn's tool still runs");
@@ -698,17 +747,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 client = MockLlmClient::with_responses(vec![tool_events, completed_text_events()]);
let tool = CountingTool::new("count_tool"); let tool = CountingTool::new("count_tool");
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
engine.register_tool(tool.definition()); engine.register_tool(tool.definition());
engine.set_interceptor(PauseToolOnce { engine.set_interceptor(PauseToolOnce {
calls: AtomicUsize::new(0), calls: AtomicUsize::new(0),
}); });
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!(engine.run("pause").await.unwrap(), EngineResult::Paused); assert_eq!(
engine.run(&mut history, "pause").await.unwrap(),
EngineResult::Paused
);
assert_eq!(engine.active_run_turn_count(), Some(1)); assert_eq!(engine.active_run_turn_count(), Some(1));
assert_eq!(engine.run("replace").await.unwrap(), EngineResult::Finished); assert_eq!(
engine.run(&mut history, "replace").await.unwrap(),
EngineResult::Finished
);
assert_eq!(engine.turn_count(), 2); assert_eq!(engine.turn_count(), 2);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(tool.call_count(), 1, "pending-tool semantics are unchanged"); assert_eq!(tool.call_count(), 1, "pending-tool semantics are unchanged");
@@ -717,14 +773,15 @@ async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() {
#[tokio::test] #[tokio::test]
async fn interceptor_continuation_consumes_the_logical_run_budget() { async fn interceptor_continuation_consumes_the_logical_run_budget() {
let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
let mut history: History = History::new();
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
engine.set_interceptor(ContinueTurnOnce { engine.set_interceptor(ContinueTurnOnce {
calls: AtomicUsize::new(0), calls: AtomicUsize::new(0),
}); });
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!( assert_eq!(
engine.run("start").await.unwrap(), engine.run(&mut history, "start").await.unwrap(),
EngineResult::LimitReached EngineResult::LimitReached
); );
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
@@ -735,13 +792,17 @@ async fn interceptor_continuation_consumes_the_logical_run_budget() {
#[tokio::test] #[tokio::test]
async fn restored_active_run_budget_is_enforced_before_another_llm_call() { async fn restored_active_run_budget_is_enforced_before_another_llm_call() {
let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
let mut history: History = History::new();
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
engine.set_turn_count(7); engine.set_turn_count(7);
engine.set_last_run_interrupted(true); engine.set_last_run_interrupted(true);
engine.set_active_run_turn_count(Some(1)); engine.set_active_run_turn_count(Some(1));
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!(engine.resume().await.unwrap(), EngineResult::LimitReached); assert_eq!(
engine.resume(&mut history).await.unwrap(),
EngineResult::LimitReached
);
assert_eq!(engine.turn_count(), 7); assert_eq!(engine.turn_count(), 7);
assert_eq!(engine.llm_call_count(), 0); assert_eq!(engine.llm_call_count(), 0);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
crates/agen/tests/parallel_execution_test.rs
+21 -9
View File
@@ -6,12 +6,12 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use agen::Engine;
use agen::interceptor::{Interceptor, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo}; use agen::interceptor::{Interceptor, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo};
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::tool::{ use agen::tool::{
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult, Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
}; };
use agen::{Engine, History};
use async_trait::async_trait; use async_trait::async_trait;
mod common; mod common;
@@ -145,6 +145,7 @@ async fn test_parallel_tool_execution() {
], ],
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let tool1 = SlowTool::new("slow_tool_1", 100); let tool1 = SlowTool::new("slow_tool_1", 100);
let tool2 = SlowTool::new("slow_tool_2", 100); let tool2 = SlowTool::new("slow_tool_2", 100);
let tool3 = SlowTool::new("slow_tool_3", 100); let tool3 = SlowTool::new("slow_tool_3", 100);
@@ -159,7 +160,7 @@ async fn test_parallel_tool_execution() {
let start = Instant::now(); let start = Instant::now();
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineResult)
let _result = engine.run("Run all tools").await; let _result = engine.run(&mut history, "Run all tools").await;
let elapsed = start.elapsed(); let elapsed = start.elapsed();
// Verify all tools were called // 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 engine = Engine::new(client);
let mut history: History = History::new();
let contexts = Arc::new(Mutex::new(Vec::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_a", contexts.clone()).definition());
engine.register_tool(ContextRecordingTool::new("record_b", contexts.clone()).definition()); engine.register_tool(ContextRecordingTool::new("record_b", contexts.clone()).definition());
engine.register_tool(ContextRecordingTool::new("record_c", 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(); let mut contexts = contexts.lock().unwrap().clone();
contexts.sort_by_key(|ctx| ctx.call_index); 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 engine = Engine::new(client);
let mut history: History = History::new();
let contexts = Arc::new(Mutex::new(Vec::new())); let contexts = Arc::new(Mutex::new(Vec::new()));
engine.register_tool(ContextRecordingTool::new("record", contexts.clone()).definition()); 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(); let contexts = contexts.lock().unwrap().clone();
assert_eq!(contexts.len(), 2); 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 engine = Engine::new(client);
let mut history: History = History::new();
let executed_contexts = Arc::new(Mutex::new(Vec::new())); let executed_contexts = Arc::new(Mutex::new(Vec::new()));
let pre_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())); 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(), 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(); let mut pre_contexts = pre_contexts.lock().unwrap().clone();
pre_contexts.sort_by_key(|ctx| ctx.call_index); 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 client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let allowed_tool = SlowTool::new("allowed_tool", 10); let allowed_tool = SlowTool::new("allowed_tool", 10);
let blocked_tool = SlowTool::new("blocked_tool", 10); let blocked_tool = SlowTool::new("blocked_tool", 10);
@@ -416,7 +423,7 @@ async fn test_before_tool_call_skip() {
engine.set_interceptor(BlockingPolicy); engine.set_interceptor(BlockingPolicy);
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineResult)
let _result = engine.run("Test hook").await; let _result = engine.run(&mut history, "Test hook").await;
// allowed_tool is called, but blocked_tool is not // allowed_tool is called, but blocked_tool is not
assert_eq!( assert_eq!(
@@ -457,6 +464,7 @@ async fn test_post_tool_call_modification() {
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
#[derive(Clone)] #[derive(Clone)]
struct SimpleTool; struct SimpleTool;
@@ -503,7 +511,7 @@ async fn test_post_tool_call_modification() {
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineResult)
let result = engine.run("Test modification").await; let result = engine.run(&mut history, "Test modification").await;
assert!(result.is_ok(), "Engine should complete"); assert!(result.is_ok(), "Engine should complete");
@@ -540,6 +548,7 @@ async fn test_before_tool_call_synthetic_result_committed() {
], ],
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let blocked_tool = SlowTool::new("blocked_tool", 10); let blocked_tool = SlowTool::new("blocked_tool", 10);
let blocked_clone = blocked_tool.clone(); let blocked_clone = blocked_tool.clone();
engine.register_tool(blocked_tool.definition()); engine.register_tool(blocked_tool.definition());
@@ -558,10 +567,13 @@ async fn test_before_tool_call_synthetic_result_committed() {
engine.set_interceptor(SyntheticPolicy); engine.set_interceptor(SyntheticPolicy);
let result = engine.run("Test synthetic result").await.unwrap(); let _result = engine
.run(&mut history, "Test synthetic result")
.await
.unwrap();
assert_eq!(blocked_clone.call_count(), 0, "Blocked tool should not run"); 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, item,
agen::Item::ToolResult { agen::Item::ToolResult {
call_id, call_id,
crates/agen/tests/reasoning_round_trip_test.rs
+26 -22
View File
@@ -13,12 +13,12 @@
mod common; mod common;
use agen::Engine;
use agen::Item; use agen::Item;
use agen::llm_client::event::{ use agen::llm_client::event::{
BlockMetadata, BlockStart, BlockStop, BlockType, Event, ReasoningBlockData, ResponseStatus, BlockMetadata, BlockStart, BlockStop, BlockType, Event, ReasoningBlockData, ResponseStatus,
StatusEvent, StatusEvent,
}; };
use agen::{Engine, History};
use common::MockLlmClient; use common::MockLlmClient;
fn reasoning_block(text: impl Into<String>, data: ReasoningBlockData) -> Vec<Event> { fn reasoning_block(text: impl Into<String>, data: ReasoningBlockData) -> Vec<Event> {
@@ -65,15 +65,15 @@ async fn anthropic_thinking_round_trips_signature_into_history() {
]); ]);
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let engine = Engine::new(client); let engine = Engine::new(client);
let out = engine.run("question?").await.expect("run ok"); let mut history: History = History::new();
let engine = out.engine; let _out = engine.run(&mut history, "question?").await.expect("run ok");
let history = engine.history(); let entries = history.entries();
// user / reasoning / assistant_message // user / reasoning / assistant_message
assert_eq!(history.len(), 3, "history: {history:?}"); assert_eq!(history.len(), 3, "history: {history:?}");
assert!(matches!(history[0], Item::Message { .. })); assert!(matches!(entries[0].item, Item::Message { .. }));
match &history[1] { match &entries[1].item {
Item::Reasoning { Item::Reasoning {
text, signature, .. text, signature, ..
} => { } => {
@@ -82,7 +82,7 @@ async fn anthropic_thinking_round_trips_signature_into_history() {
} }
other => panic!("expected Reasoning, got {other:?}"), 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 が /// 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 client = MockLlmClient::new(events);
let engine = Engine::new(client); let engine = Engine::new(client);
let out = engine.run("q").await.expect("run ok"); let mut history: History = History::new();
let engine = out.engine; let _out = engine.run(&mut history, "q").await.expect("run ok");
let history = engine.history(); let entries = history.entries();
match &history[1] { match &entries[1].item {
Item::Reasoning { Item::Reasoning {
text, text,
summary, summary,
@@ -155,13 +155,13 @@ async fn reasoning_precedes_text_in_assistant_burst() {
})); }));
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let engine = Engine::new(client); let engine = Engine::new(client);
let out = engine.run("q").await.expect("run ok"); let mut history: History = History::new();
let engine = out.engine; let _out = engine.run(&mut history, "q").await.expect("run ok");
let history = engine.history(); let entries = history.entries();
// user / reasoning(先頭) / assistant_message // user / reasoning(先頭) / assistant_message
assert!(matches!(history[1], Item::Reasoning { .. })); assert!(matches!(entries[1].item, Item::Reasoning { .. }));
assert_eq!(history[2].as_text(), Some("intermediate")); assert_eq!(entries[2].item.as_text(), Some("intermediate"));
} }
/// resume シナリオ: history.json 由来の Item::Reasoning(signature) を Engine に /// 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 engine = Engine::new(client);
let mut history: History = History::new();
// resume: 既存 history を流し込む // resume: 既存 history を流し込む
engine.set_history(vec![ engine.set_history(
Item::user_message("prior question"), &mut history,
Item::reasoning("prior thinking").with_signature("SIG-PRIOR"), vec![
Item::assistant_message("prior answer"), Item::user_message("prior question"),
]); Item::reasoning("prior thinking").with_signature("SIG-PRIOR"),
Item::assistant_message("prior answer"),
],
);
let _ = engine.run("follow up").await.expect("run ok"); let _ = engine.run(&mut history, "follow up").await.expect("run ok");
let req = captured let req = captured
.lock() .lock()
crates/agen/tests/ui/locked_register_tool.rs
+3 -2
View File
@@ -1,4 +1,4 @@
use agen::Engine; use agen::{Engine, History};
use agen::llm_client::capability::{ use agen::llm_client::capability::{
CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport, CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport,
}; };
@@ -22,7 +22,8 @@ fn main() {
cap, cap,
); );
let engine = Engine::new(client); 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 def: agen::tool::ToolDefinition = Arc::new(|| panic!("unused"));
let _ = locked.register_tool(def); let _ = locked.register_tool(def);
} }
crates/agen/tests/ui/locked_register_tool.stderr
+3 -3
View File
@@ -1,8 +1,8 @@
error[E0599]: no method named `register_tool` found for struct `Engine<HttpTransport<AnthropicScheme>, Locked>` in the current scope error[E0599]: no method named `register_tool` found for struct `Engine<HttpTransport<AnthropicScheme>, 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<HttpTransport<AnthropicScheme>, Locked>` | ^^^^^^^^^^^^^ method not found in `Engine<HttpTransport<AnthropicScheme>, Locked>`
| |
= note: the method was found for = note: the method was found for
- `Engine<C>` - `Engine<C, Mutable, A>`
crates/memory/src/extract/payload.rs
+4 -1
View File
@@ -9,7 +9,7 @@
use schemars::JsonSchema; use schemars::JsonSchema;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::schema::{EvidenceKind, SourceEvidenceRef, SourceRef}; use crate::schema::{EvidenceKind, EvidenceOrigin, SourceEvidenceRef, SourceRef};
/// Current flat staging schema version. /// Current flat staging schema version.
pub const STAGING_SCHEMA_VERSION: u32 = 2; pub const STAGING_SCHEMA_VERSION: u32 = 2;
@@ -80,6 +80,8 @@ pub struct StagingEvidence {
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub entry_range: Option<[u64; 2]>, pub entry_range: Option<[u64; 2]>,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub origin: Option<EvidenceOrigin>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub excerpt: Option<String>, pub excerpt: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub summary: Option<String>, pub summary: Option<String>,
@@ -159,6 +161,7 @@ mod tests {
id: "E001".into(), id: "E001".into(),
kind: EvidenceKind::new(EvidenceKind::MESSAGE), kind: EvidenceKind::new(EvidenceKind::MESSAGE),
entry_range: Some([10, 12]), entry_range: Some([10, 12]),
origin: None,
excerpt: Some("extract candidate taxonomy".into()), excerpt: Some("extract candidate taxonomy".into()),
summary: Some("User and assistant discussed staging kinds".into()), summary: Some("User and assistant discussed staging kinds".into()),
}; };
crates/memory/src/schema/common.rs
+37
View File
@@ -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<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub workspace_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub runtime_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub worker_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub flow_selector: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub flow_definition_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub flow_definition_revision: Option<u64>,
}
/// Host-resolved source/evidence metadata for an individual staging claim. /// Host-resolved source/evidence metadata for an individual staging claim.
/// ///
/// This deliberately stores only bounded anchor metadata: stable ids, entry /// 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. /// Host-assigned evidence id within the referenced evidence set.
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub evidence_id: Option<String>, pub evidence_id: Option<String>,
/// Trusted typed origin snapshot for this logical evidence entry.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub origin: Option<EvidenceOrigin>,
/// Extensible evidence kind tag. /// Extensible evidence kind tag.
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub evidence_kind: Option<EvidenceKind>, pub evidence_kind: Option<EvidenceKind>,
crates/memory/src/schema/mod.rs
+4 -1
View File
@@ -10,7 +10,10 @@ mod decision;
mod request; mod request;
mod summary; 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 decision::{DecisionFrontmatter, DecisionStatus};
pub use request::RequestFrontmatter; pub use request::RequestFrontmatter;
pub use summary::SummaryFrontmatter; pub use summary::SummaryFrontmatter;
crates/session-store/src/history.rs
+180
View File
@@ -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<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub runtime_id: Option<String>,
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<String>,
},
ModelOutput {
worker: LoggedWorkerSubject,
},
ToolOutput {
worker: LoggedWorkerSubject,
},
DerivedSummary,
LegacyUnknown,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LoggedHistoryDerivation {
pub sources: Vec<LoggedSessionHistoryEntryId>,
}
#[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<LoggedHistoryDerivation>,
}
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<Item = LoggedItem>,
) -> Vec<LoggedHistoryEntry> {
let _ = session_id;
items.into_iter().map(legacy_logged_history).collect()
}
crates/session-store/src/lib.rs
+6
View File
@@ -32,6 +32,7 @@
pub mod event_trace; pub mod event_trace;
pub mod fs_store; pub mod fs_store;
pub mod history;
pub mod logged_item; pub mod logged_item;
pub mod segment; pub mod segment;
pub mod segment_log; pub mod segment_log;
@@ -44,6 +45,11 @@ pub use agen::UsageRecord;
pub use agen::llm_client::types::{ContentPart, Item, Role}; pub use agen::llm_client::types::{ContentPart, Item, Role};
pub use event_trace::{TraceEntry, TracePayload}; pub use event_trace::{TraceEntry, TracePayload};
pub use fs_store::FsStore; 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 logged_item::{LoggedContentPart, LoggedItem, LoggedRole, from_logged, to_logged};
pub use segment::{ pub use segment::{
SegmentStartState, append_entry, append_system_item, classify_history_item, SegmentStartState, append_entry, append_system_item, classify_history_item,
crates/session-store/src/segment_log.rs
+77
View File
@@ -14,6 +14,7 @@ use agen::{EngineResult, UsageRecord};
use protocol::{InvokeKind, Segment}; use protocol::{InvokeKind, Segment};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::history::{LoggedHistoryEntry, LoggedSystemHistoryEntry};
use crate::logged_item::LoggedItem; use crate::logged_item::LoggedItem;
use crate::system_item::SystemItem; use crate::system_item::SystemItem;
@@ -70,6 +71,20 @@ pub enum LogEntry {
compacted_from: Option<SegmentOrigin>, compacted_from: Option<SegmentOrigin>,
}, },
/// 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<String>,
config: RequestConfig,
history: Vec<LoggedHistoryEntry>,
#[serde(default, skip_serializing_if = "Option::is_none")]
forked_from: Option<SegmentOrigin>,
#[serde(default, skip_serializing_if = "Option::is_none")]
compacted_from: Option<SegmentOrigin>,
},
/// IDLE → active marker. Records the start of a new self-driving /// IDLE → active marker. Records the start of a new self-driving
/// cycle (Invoke range). The range extends implicitly until the /// cycle (Invoke range). The range extends implicitly until the
/// next `Invoke` entry; this entry carries the trigger only — the /// next `Invoke` entry; this entry carries the trigger only — the
@@ -105,14 +120,37 @@ pub enum LogEntry {
extensions: Vec<SessionExtension>, extensions: Vec<SessionExtension>,
}, },
/// 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<Segment>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
extensions: Vec<SessionExtension>,
history: Vec<LoggedHistoryEntry>,
},
/// 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, /// One assistant-side item appended to history — assistant message,
/// reasoning, or tool call. Singular: one entry per history item so /// reasoning, or tool call. Singular: one entry per history item so
/// the wire-side `Event::*` lane and on-disk LogEntry stay 1:1. /// the wire-side `Event::*` lane and on-disk LogEntry stay 1:1.
AssistantItem { ts: u64, item: LoggedItem }, 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. /// One tool-execution result appended to history.
ToolResult { ts: u64, item: LoggedItem }, 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 /// One typed agent-injected system item: notification, child-Worker
/// lifecycle event, `@<path>` / `/<slug>` resolution payload. Each /// lifecycle event, `@<path>` / `/<slug>` resolution payload. Each
/// `SystemItem` carries kind metadata that the LLM /// `SystemItem` carries kind metadata that the LLM
@@ -278,6 +316,22 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState {
state.config = config.clone(); state.config = config.clone();
state.history = history.iter().cloned().map(Item::from).collect(); 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 { .. } => { LogEntry::Invoke { .. } => {
// A terminal run record below clears or refines this. If the // A terminal run record below clears or refines this. If the
// log ends first, restore must treat the turn as interrupted. // 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())), .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, .. } => { LogEntry::AssistantItem { item, .. } => {
state.history.push(Item::from(item.clone())); state.history.push(Item::from(item.clone()));
} }
crates/session-store/src/worker_session_store.rs
+97 -10
View File
@@ -20,7 +20,8 @@ use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::SystemTime; 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 SESSION_FILE: &str = "session.json";
const SEGMENTS_DIR: &str = "segments"; const SEGMENTS_DIR: &str = "segments";
@@ -44,15 +45,22 @@ impl WorkerSessionStore {
fs::create_dir_all(root.join(SEGMENTS_DIR))?; fs::create_dir_all(root.join(SEGMENTS_DIR))?;
let session_id = match fs::read(root.join(SESSION_FILE)) { let session_id = match fs::read(root.join(SESSION_FILE)) {
Ok(bytes) => { Ok(bytes) => {
let manifest: SessionManifest = serde_json::from_slice(&bytes)?; let mut manifest: SessionManifest = serde_json::from_slice(&bytes)?;
if manifest.schema_version != SESSION_SCHEMA_VERSION { match manifest.schema_version {
return Err(StoreError::Corrupt { SESSION_SCHEMA_VERSION => {}
line: 0, LEGACY_SESSION_SCHEMA_VERSION => {
message: format!( validate_legacy_segment_logs(&root)?;
"unsupported Worker Session schema version {}, expected {}", manifest.schema_version = SESSION_SCHEMA_VERSION;
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) 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::<LogEntry>(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<T: Serialize>(path: &Path, value: &T) -> Result<(), StoreError> { fn atomic_write_json<T: Serialize>(path: &Path, value: &T) -> Result<(), StoreError> {
let mut bytes = serde_json::to_vec_pretty(value)?; let mut bytes = serde_json::to_vec_pretty(value)?;
bytes.push(b'\n'); bytes.push(b'\n');
@@ -405,6 +444,54 @@ mod tests {
assert_eq!(store.list_sessions().unwrap(), vec![session_id]); 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] #[test]
fn reopen_preserves_session_and_segment_ids() { fn reopen_preserves_session_and_segment_ids() {
let root = tempfile::tempdir().unwrap(); let root = tempfile::tempdir().unwrap();
crates/session-store/tests/session_test.rs
+62 -28
View File
@@ -1,12 +1,13 @@
mod common; mod common;
use std::ops::{Deref, DerefMut};
use std::sync::Arc; use std::sync::Arc;
use agen::Engine;
use agen::interceptor::{Interceptor, TurnEndAction}; use agen::interceptor::{Interceptor, TurnEndAction};
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::llm_client::types::{Item, RequestConfig}; use agen::llm_client::types::{Item, RequestConfig};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, History};
use async_trait::async_trait; use async_trait::async_trait;
use common::MockLlmClient; use common::MockLlmClient;
use session_store::{FsStore, LogEntry, SegmentStartState, Store, collect_state}; use session_store::{FsStore, LogEntry, SegmentStartState, Store, collect_state};
@@ -94,15 +95,47 @@ fn make_store() -> (tempfile::TempDir, FsStore) {
(dir, store) (dir, store)
} }
struct TestWorker {
engine: Engine<MockLlmClient>,
history: History,
}
impl TestWorker {
fn new(engine: Engine<MockLlmClient>) -> Self {
Self {
engine,
history: History::new(),
}
}
fn history(&self) -> Vec<Item> {
self.history.items_cloned()
}
}
impl Deref for TestWorker {
type Target = Engine<MockLlmClient>;
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. /// Run a worker turn and persist via session-store functions.
/// Takes ownership of the worker (needed for lock/unlock) and returns it. /// Takes ownership of the worker (needed for lock/unlock) and returns it.
async fn run_and_persist( async fn run_and_persist(
worker: Engine<MockLlmClient>, mut worker: TestWorker,
store: &FsStore, store: &FsStore,
session_id: session_store::SessionId, session_id: session_store::SessionId,
segment_id: session_store::SegmentId, segment_id: session_store::SegmentId,
input: &str, input: &str,
) -> (Engine<MockLlmClient>, agen::EngineResult) { ) -> (TestWorker, agen::EngineResult) {
// Mirror Worker's run-entry contract: log the user input as segments // Mirror Worker's run-entry contract: log the user input as segments
// before the worker pushes its flattened user_message; save_delta // before the worker pushes its flattened user_message; save_delta
// skips the resulting user_message item to avoid double-write. // skips the resulting user_message item to avoid double-write.
@@ -114,13 +147,14 @@ async fn run_and_persist(
) )
.unwrap(); .unwrap();
let history_before = worker.history().len(); let history_before = worker.history.len();
let mut locked = worker.lock(); let mut locked = worker.engine.lock(&worker.history);
let result = locked.run(input).await; let result = locked.run(&mut worker.history, input).await;
let worker = locked.unlock(); 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_delta(store, session_id, segment_id, new_items).unwrap();
session_store::save_turn_end(store, session_id, segment_id, worker.turn_count()).unwrap(); session_store::save_turn_end(store, session_id, segment_id, worker.turn_count()).unwrap();
@@ -160,14 +194,14 @@ async fn run_and_persist(
async fn session_run_logs_entries() { async fn session_run_logs_entries() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(simple_text_events()); 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( let (sid, segid) = session_store::create_segment(
&store, &store,
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -204,7 +238,7 @@ async fn session_run_logs_entries() {
async fn session_restore_round_trip() { async fn session_restore_round_trip() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(simple_text_events()); 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."); worker.set_system_prompt("You are helpful.");
let (sid, segid) = session_store::create_segment( let (sid, segid) = session_store::create_segment(
@@ -212,7 +246,7 @@ async fn session_restore_round_trip() {
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -243,7 +277,7 @@ async fn session_restore_round_trip() {
async fn session_run_with_tool_call() { async fn session_run_with_tool_call() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::with_responses(tool_call_events()); 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.register_tool(weather_tool_definition());
let (sid, segid) = session_store::create_segment( let (sid, segid) = session_store::create_segment(
@@ -251,7 +285,7 @@ async fn session_run_with_tool_call() {
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -277,7 +311,7 @@ async fn session_resume_after_pause() {
// First run: tool call with pause policy → Paused // First run: tool call with pause policy → Paused
let client = MockLlmClient::with_responses(tool_call_events()); 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.register_tool(weather_tool_definition());
worker.set_interceptor(PausePolicy); worker.set_interceptor(PausePolicy);
@@ -286,7 +320,7 @@ async fn session_resume_after_pause() {
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -317,7 +351,7 @@ async fn session_resume_after_pause() {
async fn session_fork_creates_new_session() { async fn session_fork_creates_new_session() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(simple_text_events()); 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"); worker.set_system_prompt("System prompt");
let (sid, segid) = session_store::create_segment( let (sid, segid) = session_store::create_segment(
@@ -325,7 +359,7 @@ async fn session_fork_creates_new_session() {
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -338,7 +372,7 @@ async fn session_fork_creates_new_session() {
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -359,14 +393,14 @@ async fn session_fork_creates_new_session() {
async fn session_fork_at_truncates_within_session() { async fn session_fork_at_truncates_within_session() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(simple_text_events()); 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( let (sid, segid) = session_store::create_segment(
&store, &store,
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -404,14 +438,14 @@ async fn session_fork_at_truncates_within_session() {
async fn session_config_changed_logged() { async fn session_config_changed_logged() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(vec![]); 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( let (sid, segid) = session_store::create_segment(
&store, &store,
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -437,14 +471,14 @@ async fn session_auto_forks_on_conflict() {
// Create a segment // Create a segment
let client_a = MockLlmClient::new(simple_text_events()); 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( let (sid, original_segid) = session_store::create_segment(
&store, &store,
SegmentStartState { SegmentStartState {
system_prompt: worker_a.get_system_prompt(), system_prompt: worker_a.get_system_prompt(),
config: worker_a.request_config(), config: worker_a.request_config(),
history: worker_a.history(), history: &worker_a.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -470,7 +504,7 @@ async fn session_auto_forks_on_conflict() {
SegmentStartState { SegmentStartState {
system_prompt: worker_a.get_system_prompt(), system_prompt: worker_a.get_system_prompt(),
config: worker_a.request_config(), config: worker_a.request_config(),
history: worker_a.history(), history: &worker_a.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -522,14 +556,14 @@ async fn session_auto_forks_on_conflict() {
async fn nested_past_fork_leaves_ancestors_immutable() { async fn nested_past_fork_leaves_ancestors_immutable() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(simple_text_events()); 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( let (sid, root_segid) = session_store::create_segment(
&store, &store,
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
crates/worker-runtime/src/worker_backend.rs
+18 -13
View File
@@ -38,9 +38,7 @@ use crate::working_directory::{
}; };
use async_trait::async_trait; use async_trait::async_trait;
use protocol::{Event, Method, Segment, WorkerStatus}; use protocol::{Event, Method, Segment, WorkerStatus};
use session_store::{ use session_store::{CombinedStore, LogEntry, WorkerAggregateStore, WorkerSessionStore};
CombinedStore, LogEntry, WorkerAggregateStore, WorkerSessionStore, collect_state,
};
#[cfg(test)] #[cfg(test)]
use session_store::{FsStore, FsWorkerStore}; use session_store::{FsStore, FsWorkerStore};
use tokio::runtime::Runtime; 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); const USER_INPUT_COMMIT_TIMEOUT: Duration = Duration::from_secs(9);
fn user_input_has_submission(entry: &LogEntry, submission_id: &str) -> bool { fn user_input_has_submission(entry: &LogEntry, submission_id: &str) -> bool {
let LogEntry::UserInput { extensions, .. } = entry else { let extensions = match entry {
return false; LogEntry::UserInput { extensions, .. }
| LogEntry::AnnotatedUserInput { extensions, .. } => extensions,
_ => return false,
}; };
extensions.iter().any(|extension| { extensions.iter().any(|extension| {
extension.domain == WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN extension.domain == WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN
@@ -212,11 +212,11 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider {
return Err(WorkerObservationError::NotFound); return Err(WorkerObservationError::NotFound);
} }
let entries = sink.subscribe_with_snapshot().0; let entries = sink.subscribe_with_snapshot().0;
let state = collect_state(&entries); WorkerSessionCapture::from_log_entries(
Ok(WorkerSessionCapture { format!("runtime:{runtime_id}:worker:{worker_id}"),
segment_id: format!("runtime:{runtime_id}:worker:{worker_id}"), &entries,
items: state.history, )
}) .map_err(WorkerObservationError::Unavailable)
} }
} }
@@ -2505,7 +2505,9 @@ mod tests {
let scope = Scope::writable(&scope_root).map_err(|err| err.to_string())?; let scope = Scope::writable(&scope_root).map_err(|err| err.to_string())?;
let worker = Worker::new( let worker = Worker::new(
manifest, manifest,
Engine::new(self.client.clone()), Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(
self.client.clone(),
),
store, store,
workspace_context, workspace_context,
filesystem_authority, filesystem_authority,
@@ -3241,14 +3243,17 @@ mod tests {
matches!( matches!(
entry, entry,
LogEntry::UserInput { segments, .. } LogEntry::UserInput { segments, .. }
| LogEntry::AnnotatedUserInput { segments, .. }
if segments == &vec![Segment::text("start the ticket")] if segments == &vec![Segment::text("start the ticket")]
) )
})); }));
let submission_id = entries let submission_id = entries
.iter() .iter()
.find_map(|entry| { .find_map(|entry| {
let LogEntry::UserInput { extensions, .. } = entry else { let extensions = match entry {
return None; LogEntry::UserInput { extensions, .. }
| LogEntry::AnnotatedUserInput { extensions, .. } => extensions,
_ => return None,
}; };
extensions extensions
.iter() .iter()
crates/worker/examples/worker_cli.rs
+1 -1
View File
@@ -70,7 +70,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
} }
// 5. Extract the assistant's reply from history // 5. Extract the assistant's reply from history
let history = worker.engine().history(); let history = worker.history();
if let Some(text) = history if let Some(text) = history
.iter() .iter()
.rev() .rev()
crates/worker/src/compact/prune.rs
+1 -1
View File
@@ -22,7 +22,7 @@ use crate::compact::token_counter::{
EstimateSource, savings_for_prune_impl, token_estimates_for_prune_impl, EstimateSource, savings_for_prune_impl, token_estimates_for_prune_impl,
}; };
impl<C: LlmClient, St: Store> Worker<C, St> { impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// Enable prune projection on the underlying Engine. /// Enable prune projection on the underlying Engine.
/// ///
/// Registers the config and token/savings-estimator closures on the Engine. /// Registers the config and token/savings-estimator closures on the Engine.
crates/worker/src/compact/token_counter.rs
+4 -4
View File
@@ -242,13 +242,13 @@ pub(crate) fn savings_for_prune_impl(
// ── Worker に生やす公開 API ─────────────────────────────────────────────── // ── Worker に生やす公開 API ───────────────────────────────────────────────
impl<C: LlmClient, St: Store> Worker<C, St> { impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// 現在の history 全体の推定トークン数。 /// 現在の history 全体の推定トークン数。
/// ///
/// 最後の measurement と、その後に追加された未測定分の byte/4 外挿。 /// 最後の measurement と、その後に追加された未測定分の byte/4 外挿。
pub fn total_tokens(&self) -> TokenEstimate { pub fn total_tokens(&self) -> TokenEstimate {
let usage = self.usage_history(); let usage = self.usage_history();
agen::token_counter::total_tokens(self.history(), &usage) agen::token_counter::total_tokens(&self.history(), &usage)
} }
/// 任意の history index 時点でのプロンプト全長推定。 /// 任意の history index 時点でのプロンプト全長推定。
@@ -259,7 +259,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// pointer 以降に増えたプロンプト長を測るのに使う。 /// pointer 以降に増えたプロンプト長を測るのに使う。
pub fn total_tokens_at(&self, history_len: usize) -> TokenEstimate { pub fn total_tokens_at(&self, history_len: usize) -> TokenEstimate {
let usage = self.usage_history(); 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` トークン以上を残すための分割位置。 /// 末尾から `retained` トークン以上を残すための分割位置。
@@ -267,7 +267,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// `history[..cut.index]` が要約／破棄される側、`history[cut.index..]` が残る側。 /// `history[..cut.index]` が要約／破棄される側、`history[cut.index..]` が残る側。
pub fn split_for_retained(&self, retained: u64) -> SplitPoint { pub fn split_for_retained(&self, retained: u64) -> SplitPoint {
let usage = self.usage_history(); let usage = self.usage_history();
split_for_retained_impl(self.history(), &usage, retained) split_for_retained_impl(&self.history(), &usage, retained)
} }
} }
crates/worker/src/controller.rs
+3 -3
View File
@@ -1760,7 +1760,7 @@ where
fn emit_rewind_targets<C, St>(worker: &Worker<C, St>, event_tx: &broadcast::Sender<Event>) fn emit_rewind_targets<C, St>(worker: &Worker<C, St>, event_tx: &broadcast::Sender<Event>)
where where
C: LlmClient, C: LlmClient + 'static,
St: Store, St: Store,
{ {
match worker.list_rewind_targets() { match worker.list_rewind_targets() {
@@ -1786,7 +1786,7 @@ fn apply_rewind<C, St>(
expected_head_entries: usize, expected_head_entries: usize,
) -> bool ) -> bool
where where
C: LlmClient, C: LlmClient + 'static,
St: Store, St: Store,
{ {
match worker.rewind_to(target, expected_head_entries) { match worker.rewind_to(target, expected_head_entries) {
@@ -1834,7 +1834,7 @@ fn model_supports_image_attachments(model: &manifest::ModelManifest) -> bool {
fn build_greeting<C, St>(worker: &Worker<C, St>) -> protocol::Greeting fn build_greeting<C, St>(worker: &Worker<C, St>) -> protocol::Greeting
where where
C: LlmClient, C: LlmClient + 'static,
St: Store, St: Store,
{ {
let manifest = worker.manifest(); let manifest = worker.manifest();
crates/worker/src/feature.rs
+2 -2
View File
@@ -1795,9 +1795,9 @@ impl FeatureRegistryBuilder {
} }
/// Install modules into the existing Engine tool path and hook builder. /// Install modules into the existing Engine tool path and hook builder.
pub(crate) fn install_into_engine<C: LlmClient>( pub(crate) fn install_into_engine<C: LlmClient, A>(
self, self,
worker: &mut Engine<C, Mutable>, worker: &mut Engine<C, Mutable, A>,
hook_builder: &mut HookRegistryBuilder, hook_builder: &mut HookRegistryBuilder,
) -> FeatureRegistryInstallReport { ) -> FeatureRegistryInstallReport {
let mut pending_tools = Vec::new(); let mut pending_tools = Vec::new();
crates/worker/src/feature/builtin/memory_extract.rs
+98 -7
View File
@@ -6,7 +6,9 @@ use memory::backend::{
MemoryBackendOperation, MemoryBackendOperationResult, MemoryStageCandidateOperation, MemoryBackendOperation, MemoryBackendOperationResult, MemoryStageCandidateOperation,
}; };
use memory::extract::{CandidateKind, ExtractedCandidate, StagingEvidence}; use memory::extract::{CandidateKind, ExtractedCandidate, StagingEvidence};
use memory::schema::{EvidenceKind, SourceEvidenceRef, SourceRef}; use memory::schema::{
EvidenceKind, EvidenceOrigin, EvidenceOriginKind, SourceEvidenceRef, SourceRef,
};
use schemars::JsonSchema; use schemars::JsonSchema;
use serde::Deserialize; use serde::Deserialize;
@@ -174,17 +176,29 @@ impl Tool for StageMemoryCandidateTool {
"StageMemoryCandidate requires at least one entry_ref".to_string(), "StageMemoryCandidate requires at least one entry_ref".to_string(),
)); ));
} }
let mut evidence = Vec::with_capacity(params.entry_refs.len()); let mut entries = Vec::with_capacity(params.entry_refs.len());
let mut source_refs = Vec::with_capacity(params.entry_refs.len());
for entry_ref in &params.entry_refs { for entry_ref in &params.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!( ToolError::InvalidArgument(format!(
"unknown SessionEntryRef {entry_ref:?} for this extraction capture" "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 { let candidate = ExtractedCandidate {
kind: params.kind, kind: params.kind,
claim: params.claim, 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 { fn staging_evidence(entry: &SessionEntryEvidence) -> StagingEvidence {
StagingEvidence { StagingEvidence {
id: entry.entry_ref.to_string(), id: entry.entry_ref.to_string(),
kind: evidence_kind(entry), kind: evidence_kind(entry),
entry_range: Some(entry.entry_range), entry_range: Some(entry.entry_range),
origin: Some(evidence_origin(&entry.origin)),
excerpt: Some(entry.excerpt.clone()), excerpt: Some(entry.excerpt.clone()),
summary: Some(entry.summary.clone()), summary: Some(entry.summary.clone()),
} }
@@ -325,6 +393,7 @@ fn source_evidence_ref(entry: &SessionEntryEvidence) -> SourceEvidenceRef {
segment_id: Some(entry.segment_id.clone()), segment_id: Some(entry.segment_id.clone()),
entry_range: Some(entry.entry_range), entry_range: Some(entry.entry_range),
evidence_id: Some(entry.entry_ref.to_string()), evidence_id: Some(entry.entry_ref.to_string()),
origin: Some(evidence_origin(&entry.origin)),
evidence_kind: Some(evidence_kind(entry)), evidence_kind: Some(evidence_kind(entry)),
label: Some(entry.label.clone()), label: Some(entry.label.clone()),
summary: Some(entry.summary.clone()), summary: Some(entry.summary.clone()),
@@ -432,6 +501,15 @@ mod tests {
assert!(input.contains("StageMemoryCandidate.entry_refs")); 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] #[test]
fn backend_input_failures_remain_invalid_argument_tool_errors() { fn backend_input_failures_remain_invalid_argument_tool_errors() {
let backend = map_memory_stage_error(WorkspaceMemoryBackendError::Backend( let backend = map_memory_stage_error(WorkspaceMemoryBackendError::Backend(
@@ -445,6 +523,19 @@ mod tests {
assert!(matches!(http, ToolError::InvalidArgument(_))); 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] #[tokio::test]
async fn stage_rejects_entry_ref_outside_capture_before_backend_mutation() { async fn stage_rejects_entry_ref_outside_capture_before_backend_mutation() {
let tool = StageMemoryCandidateTool { state: state() }; let tool = StageMemoryCandidateTool { state: state() };
crates/worker/src/feature/builtin/session_explore.rs
+5 -2
View File
@@ -193,6 +193,7 @@ impl Tool for ShowOverviewTool {
.map(|entry| { .map(|entry| {
serde_json::json!({ serde_json::json!({
"entry_ref": entry.id, "entry_ref": entry.id,
"origin": entry.origin,
"entry_range": entry.entry_range, "entry_range": entry.entry_range,
"kind": entry.kind.as_str(), "kind": entry.kind.as_str(),
"label": entry.label, "label": entry.label,
@@ -234,15 +235,16 @@ impl Tool for SearchEntriesTool {
.transpose()?; .transpose()?;
let from = params.from.as_deref().map(parse_entry_ref).transpose()?; let from = params.from.as_deref().map(parse_entry_ref).transpose()?;
let through = params.through.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 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( return Err(ToolError::InvalidArgument(
"SearchEntries from must not be after through".to_string(), "SearchEntries from must not be after through".to_string(),
)); ));
} }
} }
let limit = bounded_limit(params.limit, DEFAULT_PAGE_LIMIT, MAX_PAGE_LIMIT); 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, query: params.query,
kind, kind,
tool_part, tool_part,
@@ -318,6 +320,7 @@ impl Tool for ReadEntryTool {
.map(|entry| { .map(|entry| {
serde_json::json!({ serde_json::json!({
"entry_ref": entry.id, "entry_ref": entry.id,
"origin": entry.origin,
"entry_range": entry.entry_range, "entry_range": entry.entry_range,
"kind": entry.kind.as_str(), "kind": entry.kind.as_str(),
"tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()), "tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()),
crates/worker/src/feature/builtin/worker_observation.rs
+66 -11
View File
@@ -1,11 +1,12 @@
use std::sync::Arc; use std::sync::Arc;
#[cfg(test)]
use agen::Item; use agen::Item;
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait; use async_trait::async_trait;
use schemars::JsonSchema; use schemars::JsonSchema;
use serde::{Deserialize, Serialize}; 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 super::manage_worker::{WORKER_CONTROL_SERVICE_ID, WorkerControlService};
use crate::feature::{ use crate::feature::{
@@ -60,7 +61,27 @@ pub struct WorkerObservationSubject {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct WorkerSessionCapture { pub struct WorkerSessionCapture {
pub segment_id: String, pub segment_id: String,
pub items: Vec<Item>, pub entries: Vec<agen::HistoryEntry<crate::SessionHistoryMetadata>>,
}
impl WorkerSessionCapture {
pub fn from_log_entries(
segment_id: impl Into<String>,
log_entries: &[LogEntry],
) -> Result<Self, String> {
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)] #[derive(Debug, thiserror::Error)]
@@ -161,9 +182,17 @@ impl WorkerObservationProvider for WorkspaceClientWorkerObservationProvider {
}) })
.collect::<Result<Vec<session_store::LogEntry>, _>>()?; .collect::<Result<Vec<session_store::LogEntry>, _>>()?;
let state = collect_state(&entries); 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 { Ok(WorkerSessionCapture {
segment_id: response.segment_id, segment_id,
items: state.history, entries: typed_entries,
}) })
} }
} }
@@ -392,9 +421,15 @@ impl WorkerObservationProvider for SpawnedSubWorkerObservationProvider {
.ok_or(WorkerObservationError::NotFound)?; .ok_or(WorkerObservationError::NotFound)?;
let entries = record.session.entries(); let entries = record.session.entries();
let state = collect_state(&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 { Ok(WorkerSessionCapture {
segment_id: format!("subworker:{name}"), segment_id: format!("subworker:{name}"),
items: state.history, entries: typed_entries,
}) })
} }
} }
@@ -508,6 +543,7 @@ impl Tool for ViewSessionOverviewTool {
.map(|entry| { .map(|entry| {
serde_json::json!({ serde_json::json!({
"entry_ref": entry.id, "entry_ref": entry.id,
"origin": entry.origin,
"entry_range": entry.entry_range, "entry_range": entry.entry_range,
"kind": entry.kind.as_str(), "kind": entry.kind.as_str(),
"label": entry.label, "label": entry.label,
@@ -547,7 +583,7 @@ impl Tool for SearchSessionEntriesTool {
let from = params.from.as_deref().map(parse_entry_ref).transpose()?; let from = params.from.as_deref().map(parse_entry_ref).transpose()?;
let through = params.through.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 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( return Err(ToolError::InvalidArgument(
"SearchSessionEntries from must not be after through".to_string(), "SearchSessionEntries from must not be after through".to_string(),
)); ));
@@ -573,6 +609,7 @@ impl Tool for SearchSessionEntriesTool {
.map(|entry| { .map(|entry| {
serde_json::json!({ serde_json::json!({
"entry_ref": entry.id, "entry_ref": entry.id,
"origin": entry.origin,
"entry_range": entry.entry_range, "entry_range": entry.entry_range,
"kind": entry.kind.as_str(), "kind": entry.kind.as_str(),
"tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()), "tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()),
@@ -628,6 +665,7 @@ impl Tool for ReadSessionEntryTool {
.map(|entry| { .map(|entry| {
serde_json::json!({ serde_json::json!({
"entry_ref": entry.id, "entry_ref": entry.id,
"origin": entry.origin,
"entry_range": entry.entry_range, "entry_range": entry.entry_range,
"kind": entry.kind.as_str(), "kind": entry.kind.as_str(),
"tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()), "tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()),
@@ -661,7 +699,10 @@ async fn latest_view(
.capture_worker_session(subject) .capture_worker_session(subject)
.await .await
.map_err(tool_error)?; .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<T: serde::de::DeserializeOwned>( fn parse_input<T: serde::de::DeserializeOwned>(
@@ -751,9 +792,23 @@ mod tests {
if subject != &granted_subject() { if subject != &granted_subject() {
return Err(WorkerObservationError::NotFound); 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 { Ok(WorkerSessionCapture {
segment_id: "segment".to_string(), 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 read = read_definition(provider.clone())().1;
let hidden = read let hidden = read
.execute( .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(), agen::tool::ToolExecutionContext::direct(),
) )
.await .await
@@ -810,7 +865,7 @@ mod tests {
.push(message("a1", Role::Assistant, "second")); .push(message("a1", Role::Assistant, "second"));
let output = read let output = read
.execute( .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(), agen::tool::ToolExecutionContext::direct(),
) )
.await .await
@@ -819,7 +874,7 @@ mod tests {
let output = read let output = read
.execute( .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(), agen::tool::ToolExecutionContext::direct(),
) )
.await .await
crates/worker/src/internal_worker.rs
+20 -4
View File
@@ -55,7 +55,17 @@ pub(crate) struct InternalWorkerSpec {
pub input: String, pub input: String,
pub cache_key: Option<String>, pub cache_key: Option<String>,
pub max_turns: Option<u32>, pub max_turns: Option<u32>,
pub engine_configurator: Option<Box<dyn FnOnce(&mut Engine<Box<dyn LlmClient>>) + Send>>, pub engine_configurator: Option<
Box<
dyn FnOnce(
&mut Engine<
Box<dyn LlmClient>,
agen::state::Mutable,
crate::SessionHistoryMetadata,
>,
) + Send,
>,
>,
pub features: FeatureRegistryBuilder, pub features: FeatureRegistryBuilder,
pub required_tools: &'static [&'static str], pub required_tools: &'static [&'static str],
pub authority: InternalWorkerAuthority, pub authority: InternalWorkerAuthority,
@@ -124,7 +134,9 @@ where
let last_usage = Arc::new(Mutex::new(None::<UsageEvent>)); let last_usage = Arc::new(Mutex::new(None::<UsageEvent>));
let usage_slot = last_usage.clone(); 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| { engine.on_usage(move |usage| {
if let Ok(mut slot) = usage_slot.lock() { if let Ok(mut slot) = usage_slot.lock() {
*slot = Some(usage.clone()); *slot = Some(usage.clone());
@@ -494,7 +506,9 @@ pub(crate) async fn spawn_internal_worker_session(
let last_usage = Arc::new(Mutex::new(None::<UsageEvent>)); let last_usage = Arc::new(Mutex::new(None::<UsageEvent>));
let usage_slot = last_usage.clone(); 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| { engine.on_usage(move |usage| {
if let Ok(mut slot) = usage_slot.lock() { if let Ok(mut slot) = usage_slot.lock() {
*slot = Some(usage.clone()); *slot = Some(usage.clone());
@@ -591,7 +605,9 @@ pub(crate) fn prepare_internal_worker_from_spec(
manifest.compaction = None; manifest.compaction = None;
manifest.memory = 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_cache_key(cache_key);
engine.set_max_turns(max_turns); engine.set_max_turns(max_turns);
if let Some(configure) = engine_configurator { if let Some(configure) = engine_configurator {
crates/worker/src/ipc/interceptor.rs
+39 -2
View File
@@ -8,6 +8,7 @@
//! decisions (continue / skip / abort / pause). //! decisions (continue / skip / abort / pause).
use std::borrow::Cow; use std::borrow::Cow;
use std::collections::VecDeque;
use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex}; 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::ipc::notify_buffer::{NotifyBuffer, build_system_item_with_provenance};
use crate::prompt::catalog::PromptCatalog; use crate::prompt::catalog::PromptCatalog;
use crate::session_history::SessionHistoryMetadata;
use crate::worker::SystemItemCommitter; use crate::worker::SystemItemCommitter;
use agen::HistoryEntry;
use agen::token_counter::total_tokens; use agen::token_counter::total_tokens;
/// Maximum number of bytes copied into `TurnEndInfo::final_text_preview`. /// 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 /// worker. `None` in tests / `Worker::new` paths where no writer is
/// attached. /// attached.
log_writer: Option<Arc<dyn SystemItemCommitter>>, log_writer: Option<Arc<dyn SystemItemCommitter>>,
pending_committed_history: Arc<Mutex<VecDeque<HistoryEntry<SessionHistoryMetadata>>>>,
/// Next turn index assigned by `on_prompt_submit`. /// Next turn index assigned by `on_prompt_submit`.
next_turn_index: AtomicUsize, next_turn_index: AtomicUsize,
/// Tool calls observed in the current turn (reset on each new prompt). /// Tool calls observed in the current turn (reset on each new prompt).
@@ -80,6 +84,7 @@ pub(crate) struct WorkerInterceptor {
} }
impl WorkerInterceptor { impl WorkerInterceptor {
#[cfg(test)]
pub(crate) fn new( pub(crate) fn new(
registry: Arc<HookRegistry>, registry: Arc<HookRegistry>,
compact_state: Option<Arc<CompactState>>, compact_state: Option<Arc<CompactState>>,
@@ -88,6 +93,28 @@ impl WorkerInterceptor {
pending_attachments: Arc<Mutex<Vec<SystemItem>>>, pending_attachments: Arc<Mutex<Vec<SystemItem>>>,
prompts: Arc<ArcSwap<PromptCatalog>>, prompts: Arc<ArcSwap<PromptCatalog>>,
log_writer: Option<Arc<dyn SystemItemCommitter>>, log_writer: Option<Arc<dyn SystemItemCommitter>>,
) -> 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<HookRegistry>,
compact_state: Option<Arc<CompactState>>,
usage_history: Option<Arc<Mutex<Vec<UsageRecord>>>>,
pending_notifies: NotifyBuffer,
pending_attachments: Arc<Mutex<Vec<SystemItem>>>,
prompts: Arc<ArcSwap<PromptCatalog>>,
log_writer: Option<Arc<dyn SystemItemCommitter>>,
pending_committed_history: Arc<Mutex<VecDeque<HistoryEntry<SessionHistoryMetadata>>>>,
) -> Self { ) -> Self {
Self { Self {
registry, registry,
@@ -99,6 +126,7 @@ impl WorkerInterceptor {
prompts, prompts,
prompt_workspace_id: None, prompt_workspace_id: None,
log_writer, log_writer,
pending_committed_history,
next_turn_index: AtomicUsize::new(0), next_turn_index: AtomicUsize::new(0),
tool_calls_this_turn: AtomicUsize::new(0), tool_calls_this_turn: AtomicUsize::new(0),
} }
@@ -125,7 +153,11 @@ impl WorkerInterceptor {
return Ok(()); return Ok(());
}; };
for item in items { 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(()) Ok(())
} }
@@ -507,7 +539,12 @@ mod tests {
&self, &self,
entry: session_store::LogEntry, entry: session_store::LogEntry,
) -> Result<(), session_store::StoreError> { ) -> 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 self.committed
.lock() .lock()
.expect("committed system-item list poisoned") .expect("committed system-item list poisoned")
crates/worker/src/ipc/protocol_session.rs
+8 -2
View File
@@ -29,15 +29,21 @@ pub fn subscribe_worker_protocol_session(handle: &WorkerHandle) -> WorkerProtoco
pub fn live_log_entry_event(entry: LogEntry) -> Option<Event> { pub fn live_log_entry_event(entry: LogEntry) -> Option<Event> {
match entry { match entry {
LogEntry::SegmentStart { .. } => { entry @ (LogEntry::SegmentStart { .. } | LogEntry::AnnotatedSegmentStart { .. }) => {
let value = serde_json::to_value(&entry).expect("LogEntry is Serialize"); let value = serde_json::to_value(&entry).expect("LogEntry is Serialize");
Some(Event::SegmentRotated { entry: value }) 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, .. } => { LogEntry::SystemItem { item, .. } => {
let value = serde_json::to_value(&item).expect("SystemItem is Serialize"); let value = serde_json::to_value(&item).expect("SystemItem is Serialize");
Some(Event::SystemItem { item: value }) 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 }), LogEntry::Invoke { trigger, .. } => Some(Event::InvokeStart { kind: trigger }),
other => { other => {
// `SegmentLogSink::is_live_relevant` keeps non-live-relevant // `SegmentLogSink::is_live_relevant` keeps non-live-relevant
crates/worker/src/lib.rs
+5
View File
@@ -12,6 +12,7 @@ pub mod prompt;
pub mod runtime; pub mod runtime;
pub mod segment_log_sink; pub mod segment_log_sink;
mod session_capture; mod session_capture;
mod session_history;
pub mod shared_state; pub mod shared_state;
mod shutdown_after_idle; mod shutdown_after_idle;
pub mod skill; pub mod skill;
@@ -41,6 +42,10 @@ pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTem
pub use protocol::{ErrorCode, Event, Method, TurnResult, WorkerStatus}; pub use protocol::{ErrorCode, Event, Method, TurnResult, WorkerStatus};
pub use runtime::dir::RuntimeDir; pub use runtime::dir::RuntimeDir;
pub use segment_log_sink::SegmentLogSink; pub use segment_log_sink::SegmentLogSink;
pub use session_history::{
SessionHistoryDerivation, SessionHistoryEntryId, SessionHistoryMetadata,
WorkerHistoryProvenance, WorkerSubjectSnapshot,
};
pub use shared_state::WorkerSharedState; pub use shared_state::WorkerSharedState;
pub use worker::{ pub use worker::{
LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError, LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError,
crates/worker/src/permission.rs
+1 -1
View File
@@ -34,7 +34,7 @@ impl PermissionHook {
} }
} }
impl<C: LlmClient, St: Store> Worker<C, St> { impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
pub(crate) fn apply_permissions_from_manifest(&mut self) { pub(crate) fn apply_permissions_from_manifest(&mut self) {
let Some(permissions) = self.manifest().permissions.clone() else { let Some(permissions) = self.manifest().permissions.clone() else {
return; return;
crates/worker/src/segment_log_sink.rs
+3
View File
@@ -121,8 +121,11 @@ impl SegmentLogSink {
matches!( matches!(
entry, entry,
LogEntry::SegmentStart { .. } LogEntry::SegmentStart { .. }
| LogEntry::AnnotatedSegmentStart { .. }
| LogEntry::UserInput { .. } | LogEntry::UserInput { .. }
| LogEntry::AnnotatedUserInput { .. }
| LogEntry::SystemItem { .. } | LogEntry::SystemItem { .. }
| LogEntry::AnnotatedSystemItem { .. }
| LogEntry::Invoke { .. } | LogEntry::Invoke { .. }
) )
} }
crates/worker/src/session_capture.rs
+144 -20
View File
@@ -6,7 +6,8 @@
use std::sync::Arc; use std::sync::Arc;
use agen::{Item, Role}; use crate::session_history::{SessionHistoryMetadata, WorkerHistoryProvenance};
use agen::{HistoryEntry, Item, Role};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
const DEFAULT_SEARCH_LIMIT: usize = 20; const DEFAULT_SEARCH_LIMIT: usize = 20;
@@ -21,14 +22,21 @@ const OVERVIEW_ANCHOR_STRIDE: usize = 8;
pub(crate) struct SessionEntryRef(String); pub(crate) struct SessionEntryRef(String);
impl SessionEntryRef { impl SessionEntryRef {
pub(crate) fn new(source_index: usize) -> Self { pub(crate) fn from_history_entry_id(entry_id: &crate::SessionHistoryEntryId) -> Self {
Self(format!("E{source_index:08}")) Self(format!("E{}", entry_id.0))
} }
pub(crate) fn parse(value: &str) -> Option<Self> { pub(crate) fn parse(value: &str) -> Option<Self> {
let reference = Self(value.to_string()); let suffix = value.strip_prefix('E')?;
reference.source_index()?; if suffix.is_empty()
Some(reference) || 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 { pub(crate) fn as_str(&self) -> &str {
@@ -97,6 +105,7 @@ impl ToolPart {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct OverviewItem { pub(crate) struct OverviewItem {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance,
pub entry_range: [u64; 2], pub entry_range: [u64; 2],
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub label: String, pub label: String,
@@ -107,6 +116,7 @@ pub(crate) struct OverviewItem {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct ReferenceEntry { pub(crate) struct ReferenceEntry {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance,
pub entry_range: [u64; 2], pub entry_range: [u64; 2],
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
@@ -132,6 +142,7 @@ pub(crate) struct SearchOptions {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct SearchHit { pub(crate) struct SearchHit {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance,
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
pub tool_name: Option<String>, pub tool_name: Option<String>,
@@ -177,6 +188,7 @@ impl Default for ReadOptions {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct ReadEntry { pub(crate) struct ReadEntry {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance,
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
pub tool_name: Option<String>, pub tool_name: Option<String>,
@@ -195,6 +207,7 @@ pub(crate) struct ReadResult {
pub(crate) struct SessionEntryEvidence { pub(crate) struct SessionEntryEvidence {
pub segment_id: String, pub segment_id: String,
pub entry_ref: SessionEntryRef, pub entry_ref: SessionEntryRef,
pub origin: WorkerHistoryProvenance,
pub entry_range: [u64; 2], pub entry_range: [u64; 2],
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
@@ -206,26 +219,42 @@ pub(crate) struct SessionEntryEvidence {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct SessionCapture { pub(crate) struct SessionCapture {
segment_id: String, segment_id: String,
items: Arc<Vec<Item>>, entries: Arc<Vec<HistoryEntry<SessionHistoryMetadata>>>,
overview: Vec<OverviewItem>, overview: Vec<OverviewItem>,
index: Vec<ReferenceEntry>, index: Vec<ReferenceEntry>,
} }
impl SessionCapture { impl SessionCapture {
pub(crate) fn new(segment_id: impl Into<String>, items: Vec<Item>) -> Self { pub(crate) fn new(segment_id: impl Into<String>, items: Vec<Item>) -> 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<String>,
entries: Vec<HistoryEntry<SessionHistoryMetadata>>,
) -> Self {
let segment_id = segment_id.into(); let segment_id = segment_id.into();
let items = Arc::new(items); let entries = Arc::new(entries);
let mut overview = Vec::new(); let mut overview = Vec::new();
let mut index = 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]; let entry_range = [idx as u64, idx as u64];
match item { match item {
Item::Message { role, content, .. } => { Item::Message { role, content, .. } => {
let kind = match role { let Some(kind) = message_reference_kind(&entry.annotation.origin, role) else {
Role::User => ReferenceKind::User, continue;
Role::Assistant => ReferenceKind::Assistant,
Role::System => continue,
}; };
let text = content let text = content
.iter() .iter()
@@ -234,9 +263,10 @@ impl SessionCapture {
.join(""); .join("");
let label = format!("{} message", kind.as_str()); let label = format!("{} message", kind.as_str());
let summary = truncate_chars(&text, 240); 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 { index.push(ReferenceEntry {
id: id.clone(), id: id.clone(),
origin: entry.annotation.origin.clone(),
entry_range, entry_range,
kind, kind,
tool_part: None, tool_part: None,
@@ -248,6 +278,7 @@ impl SessionCapture {
if matches!(kind, ReferenceKind::User | ReferenceKind::Assistant) { if matches!(kind, ReferenceKind::User | ReferenceKind::Assistant) {
overview.push(OverviewItem { overview.push(OverviewItem {
id: id.clone(), id: id.clone(),
origin: entry.annotation.origin.clone(),
entry_range, entry_range,
kind, kind,
label, label,
@@ -261,7 +292,8 @@ impl SessionCapture {
} => { } => {
let text = format!("{name}\n{arguments}"); let text = format!("{name}\n{arguments}");
index.push(ReferenceEntry { index.push(ReferenceEntry {
id: SessionEntryRef::new(idx), id: SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id),
origin: entry.annotation.origin.clone(),
entry_range, entry_range,
kind: ReferenceKind::Tool, kind: ReferenceKind::Tool,
tool_part: Some(ToolPart::Input), tool_part: Some(ToolPart::Input),
@@ -287,7 +319,8 @@ impl SessionCapture {
content.as_deref().unwrap_or_default(), content.as_deref().unwrap_or_default(),
); );
index.push(ReferenceEntry { index.push(ReferenceEntry {
id: SessionEntryRef::new(idx), id: SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id),
origin: entry.annotation.origin.clone(),
entry_range, entry_range,
kind: ReferenceKind::Tool, kind: ReferenceKind::Tool,
tool_part: Some(ToolPart::Output), tool_part: Some(ToolPart::Output),
@@ -327,7 +360,7 @@ impl SessionCapture {
Self { Self {
segment_id, segment_id,
items, entries,
overview, overview,
index, index,
} }
@@ -337,6 +370,14 @@ impl SessionCapture {
&self.overview &self.overview
} }
pub(crate) fn source_index_for_ref(&self, reference: &SessionEntryRef) -> Option<u64> {
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<SearchHit> { pub(crate) fn search(&self, options: &SearchOptions) -> Vec<SearchHit> {
let query = options.query.trim().to_lowercase(); let query = options.query.trim().to_lowercase();
let limit = options let limit = options
@@ -347,12 +388,12 @@ impl SessionCapture {
let min_entry_index = options let min_entry_index = options
.from .from
.as_ref() .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)); .unwrap_or_else(|| options.min_entry_index.unwrap_or(0));
let max_entry_index = options let max_entry_index = options
.through .through
.as_ref() .as_ref()
.and_then(SessionEntryRef::source_index) .and_then(|reference| self.source_index_for_ref(reference))
.unwrap_or(u64::MAX); .unwrap_or(u64::MAX);
let mut skipped = 0usize; let mut skipped = 0usize;
let mut hits = Vec::new(); let mut hits = Vec::new();
@@ -391,6 +432,7 @@ impl SessionCapture {
} }
hits.push(SearchHit { hits.push(SearchHit {
id: entry.id.clone(), id: entry.id.clone(),
origin: entry.origin.clone(),
kind: entry.kind, kind: entry.kind,
tool_part: entry.tool_part, tool_part: entry.tool_part,
tool_name: entry.tool_name.clone(), 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; continue;
}; };
let text = render_item(item, entry, options.detail, max_bytes.saturating_sub(bytes)); let text = render_item(item, entry, options.detail, max_bytes.saturating_sub(bytes));
bytes = bytes.saturating_add(text.len()); bytes = bytes.saturating_add(text.len());
entries.push(ReadEntry { entries.push(ReadEntry {
id: entry.id.clone(), id: entry.id.clone(),
origin: entry.origin.clone(),
kind: entry.kind, kind: entry.kind,
tool_part: entry.tool_part, tool_part: entry.tool_part,
tool_name: entry.tool_name.clone(), tool_name: entry.tool_name.clone(),
@@ -485,6 +532,7 @@ impl SessionCapture {
Some(SessionEntryEvidence { Some(SessionEntryEvidence {
segment_id: self.segment_id.clone(), segment_id: self.segment_id.clone(),
entry_ref: entry.id.clone(), entry_ref: entry.id.clone(),
origin: entry.origin.clone(),
entry_range: entry.entry_range, entry_range: entry.entry_range,
kind: entry.kind, kind: entry.kind,
tool_part: entry.tool_part, tool_part: entry.tool_part,
@@ -495,6 +543,28 @@ impl SessionCapture {
} }
} }
fn message_reference_kind(
origin: &WorkerHistoryProvenance,
provider_role: &Role,
) -> Option<ReferenceKind> {
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( fn render_item(
item: &Item, item: &Item,
entry: &ReferenceEntry, entry: &ReferenceEntry,
@@ -563,6 +633,60 @@ fn truncate_chars(text: &str, max_chars: usize) -> String {
mod tests { mod tests {
use super::*; 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] #[test]
fn overview_contains_user_and_assistant_only() { fn overview_contains_user_and_assistant_only() {
let view = SessionCapture::new( let view = SessionCapture::new(
crates/worker/src/session_history.rs
+219
View File
@@ -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<SessionHistoryDerivation>,
) -> SessionHistoryMetadata {
SessionHistoryMetadata {
entry_id: SessionHistoryEntryId::new(),
origin,
derivation,
}
}
pub(crate) fn history_entry(
item: Item,
origin: WorkerHistoryProvenance,
) -> HistoryEntry<SessionHistoryMetadata> {
HistoryEntry::new(item, metadata(origin, None))
}
pub(crate) fn to_logged_history_entry(
entry: &HistoryEntry<SessionHistoryMetadata>,
) -> LoggedHistoryEntry {
LoggedHistoryEntry {
item: entry.item.clone().into(),
metadata: entry.annotation.clone(),
}
}
fn legacy_entry(item: Item) -> HistoryEntry<SessionHistoryMetadata> {
HistoryEntry::new(item, SessionHistoryMetadata::legacy_unknown())
}
fn from_logged(entry: &LoggedHistoryEntry) -> HistoryEntry<SessionHistoryMetadata> {
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<Vec<HistoryEntry<SessionHistoryMetadata>>, 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]
);
}
}
crates/worker/src/spawn/tool.rs
+2 -2
View File
@@ -1249,7 +1249,7 @@ extract_threshold = 4000
) )
.await .await
.unwrap(); .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")))) matches!(item, Item::Message { role: Role::Assistant, content, .. } if content.iter().any(|part| matches!(part, ContentPart::Text { text } if text.contains("reviewed"))))
})); }));
@@ -1271,7 +1271,7 @@ extract_threshold = 4000
) )
.await .await
.unwrap(); .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); fail_requests.store(true, Ordering::SeqCst);
send.execute( send.execute(
crates/worker/src/worker.rs
+605 -204
View File
File diff suppressed because it is too large Load Diff
crates/worker/tests/compact_events_test.rs
+35 -23
View File
@@ -163,7 +163,8 @@ async fn make_worker_with_manifest(
let scope = worker::Scope::writable(&pwd).unwrap(); let scope = worker::Scope::writable(&pwd).unwrap();
std::mem::forget(pwd_tmp); 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( let mut worker = Worker::new(
manifest, manifest,
worker, worker,
@@ -204,28 +205,34 @@ fn system_texts_in_sink_session_start(
) -> Vec<String> { ) -> Vec<String> {
let (entries, _rx) = worker.sink().subscribe_with_snapshot(); let (entries, _rx) = worker.sink().subscribe_with_snapshot();
for entry in entries.into_iter().rev() { for entry in entries.into_iter().rev() {
if let session_store::LogEntry::SegmentStart { history, .. } = entry { let history = match entry {
return history session_store::LogEntry::AnnotatedSegmentStart { history, .. } => history
.into_iter() .into_iter()
.filter_map(|logged| { .map(|entry| entry.item)
let item: Item = logged.into(); .collect::<Vec<_>>(),
match item { session_store::LogEntry::SegmentStart { history, .. } => history,
Item::Message { _ => continue,
role: agen::Role::System, };
content, return history
.. .into_iter()
} => Some( .filter_map(|logged| {
content let item: Item = logged.into();
.iter() match item {
.map(|p| p.as_text().to_owned()) Item::Message {
.collect::<Vec<_>>() role: agen::Role::System,
.join(""), content,
), ..
_ => None, } => Some(
} content
}) .iter()
.collect(); .map(|p| p.as_text().to_owned())
} .collect::<Vec<_>>()
.join(""),
),
_ => None,
}
})
.collect();
} }
Vec::new() Vec::new()
} }
@@ -337,7 +344,12 @@ permission = "write"
// New segment records forked_from pointing at the source. // New segment records forked_from pointing at the source.
let new_entries = store.read_all(session_id, new_segment_id).unwrap(); let new_entries = store.read_all(session_id, new_segment_id).unwrap();
match &new_entries[0] { match &new_entries[0] {
LogEntry::SegmentStart { LogEntry::AnnotatedSegmentStart {
session_id: seg_session,
forked_from: Some(origin),
..
}
| LogEntry::SegmentStart {
session_id: seg_session, session_id: seg_session,
forked_from: Some(origin), forked_from: Some(origin),
.. ..
crates/worker/tests/controller_test.rs
+45 -26
View File
@@ -32,16 +32,29 @@ fn history_from_sink(handle: &WorkerHandle) -> Vec<Item> {
let mut items = Vec::new(); let mut items = Vec::new();
for entry in entries { for entry in entries {
match entry { match entry {
LogEntry::AnnotatedSegmentStart { history, .. } => {
items.extend(history.into_iter().map(|entry| Item::from(entry.item)));
}
LogEntry::SegmentStart { history, .. } => { LogEntry::SegmentStart { history, .. } => {
items.extend(history.into_iter().map(Item::from)); 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, .. } => { LogEntry::UserInput { segments, .. } => {
let text = protocol::Segment::flatten_to_text(&segments); let text = protocol::Segment::flatten_to_text(&segments);
items.push(Item::user_message(text)); items.push(Item::user_message(text));
} }
LogEntry::AnnotatedAssistantItem { entry, .. }
| LogEntry::AnnotatedToolResult { entry, .. } => {
items.push(Item::from(entry.item));
}
LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => { LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => {
items.push(Item::from(item)); items.push(Item::from(item));
} }
LogEntry::AnnotatedSystemItem { entry, .. } => {
items.push(entry.item.to_history_item());
}
LogEntry::SystemItem { item, .. } => { LogEntry::SystemItem { item, .. } => {
items.push(item.to_history_item()); items.push(item.to_history_item());
} }
@@ -51,6 +64,14 @@ fn history_from_sink(handle: &WorkerHandle) -> Vec<Item> {
items 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 // Mock LLM Client
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
@@ -192,7 +213,8 @@ async fn make_worker_with_pwd_and_manifest(
let scope = manifest::Scope::writable(&pwd).unwrap(); let scope = manifest::Scope::writable(&pwd).unwrap();
std::mem::forget(pwd_tmp); 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 authority = WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone());
let worker = Worker::new( let worker = Worker::new(
manifest, manifest,
@@ -804,10 +826,12 @@ async fn snapshot_includes_user_input_for_in_flight_turn() {
// Walk the entries, find a `LogEntry::UserInput` and // Walk the entries, find a `LogEntry::UserInput` and
// confirm its segments flatten to our submitted text. // confirm its segments flatten to our submitted text.
let mut found = false; let mut found = false;
for value in entries { for value in &entries {
let entry: session_store::LogEntry = let entry: session_store::LogEntry =
serde_json::from_value(value).expect("LogEntry deserialise"); serde_json::from_value(value.clone()).expect("LogEntry deserialise");
if let session_store::LogEntry::UserInput { segments, .. } = entry { if let session_store::LogEntry::UserInput { segments, .. }
| session_store::LogEntry::AnnotatedUserInput { segments, .. } = entry
{
let text = protocol::Segment::flatten_to_text(&segments); let text = protocol::Segment::flatten_to_text(&segments);
if text == "hello in-flight" { if text == "hello in-flight" {
found = true; 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; return;
} }
Event::Alert(_) => continue, 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 { 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); user_input_segments = Some(segments);
if saw_turn_end { if saw_turn_end {
break; 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 (entries, _) = handle.sink.subscribe_with_snapshot();
let saw_notify_in_mirror = entries.iter().any(|e| { let saw_notify_in_mirror = entries.iter().any(|e| {
matches!( matches!(
e, system_item(e),
session_store::LogEntry::SystemItem { Some(session_store::SystemItem::Notification { message, .. }) if message == "turn finished"
item: session_store::SystemItem::Notification { message, .. },
..
} if message == "turn finished"
) )
}); });
assert!( 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 (entries, _) = handle.sink.subscribe_with_snapshot();
let saw_worker_event_in_mirror = entries.iter().any(|e| { let saw_worker_event_in_mirror = entries.iter().any(|e| {
matches!( matches!(
e, system_item(e),
session_store::LogEntry::SystemItem { Some(session_store::SystemItem::WorkerEvent {
item: session_store::SystemItem::WorkerEvent { event: protocol::WorkerEvent::TurnEnded { worker_name },
event: protocol::WorkerEvent::TurnEnded { worker_name },
..
},
.. ..
} if worker_name == "child" }) if worker_name == "child"
) )
}); });
assert!( 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 (entries, _) = handle.sink.subscribe_with_snapshot();
let saw_scope_event_in_mirror = entries.iter().any(|entry| { let saw_scope_event_in_mirror = entries.iter().any(|entry| {
matches!( matches!(
entry, system_item(entry),
session_store::LogEntry::SystemItem { Some(session_store::SystemItem::WorkerEvent {
item: session_store::SystemItem::WorkerEvent { event: protocol::WorkerEvent::ScopeSubDelegated { .. },
event: protocol::WorkerEvent::ScopeSubDelegated { .. },
..
},
.. ..
} })
) )
}); });
assert!( assert!(
@@ -2373,7 +2391,8 @@ async fn snapshot_contains_user_input(handle: &WorkerHandle, needle: &str) -> bo
let entry: session_store::LogEntry = let entry: session_store::LogEntry =
serde_json::from_value(value).expect("LogEntry deserialise"); serde_json::from_value(value).expect("LogEntry deserialise");
match entry { 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) protocol::Segment::flatten_to_text(&segments).contains(needle)
} }
_ => false, _ => false,
crates/worker/tests/session_metrics_test.rs
+6 -3
View File
@@ -188,7 +188,8 @@ async fn make_worker(
let pwd = pwd_tmp.path().to_path_buf(); let pwd = pwd_tmp.path().to_path_buf();
let scope = worker::Scope::writable(&pwd).unwrap(); 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)); worker.register_tool(big_content_tool_definition(tool_name));
let worker = Worker::new( 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 // protected token budget covers the only user message). That is enough to drive
// the failure path: at least one metric attempts to write. // the failure path: at least one metric attempts to write.
let client = MockClient::new(vec![text_response_with_cache("hi", 0, 0)]); 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( let mut worker = Worker::new(
manifest, manifest,
worker, worker,
@@ -536,7 +538,8 @@ permission = "write"
let pwd_tmp = tempfile::tempdir().unwrap(); let pwd_tmp = tempfile::tempdir().unwrap();
let pwd = pwd_tmp.path().to_path_buf(); let pwd = pwd_tmp.path().to_path_buf();
let scope = worker::Scope::writable(&pwd).unwrap(); 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( let mut worker = Worker::new(
manifest, manifest,
worker, worker,
crates/worker/tests/system_prompt_template_test.rs
+2 -1
View File
@@ -130,7 +130,8 @@ async fn make_worker_with_body(
EffectivePromptCatalog::new(templates, 1, "test-schema", "test-toolchain").unwrap(); EffectivePromptCatalog::new(templates, 1, "test-schema", "test-toolchain").unwrap();
let loader = PromptCatalogSource::builtins_only().with_effective_catalog(projection); 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( let mut worker = Worker::new(
manifest, manifest,
worker, worker,