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

fix: reconcile typed history with run exits

Browse Source
This commit is contained in:
Hare
2026-08-27 15:12:05 +09:00
parent ec798c58d7
commit 1515a2fb86
11 changed files with 144 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/README.md
+2 -2
View File
@@ -32,10 +32,10 @@ async fn conversation<C: LlmClient>(client: C) -> Result<(), EngineError> {
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(&mut history, "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(&mut history, "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
+9 -11
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, History}; use agen::{Engine, EngineRunExit, StopReason};
use std::time::Duration; use std::time::Duration;
#[tokio::main] #[tokio::main]
@@ -29,7 +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(); let mut history = agen::History::new();
println!("🚀 Starting Engine..."); println!("🚀 Starting Engine...");
println!("💡 Will cancel after 2 seconds\n"); println!("💡 Will cancel after 2 seconds\n");
@@ -46,15 +46,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
println!("📡 Sending request to LLM..."); println!("📡 Sending request to LLM...");
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 { let output = engine.run(&mut history, "Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await;
Ok(out) => match out.result { match output.result {
EngineResult::Finished => println!("✅ Task completed normally"), EngineRunExit::Finished => println!("✅ Task completed normally"),
EngineResult::Paused => println!("⏸️ Task paused"), EngineRunExit::Paused => println!("⏸️ Task paused"),
EngineResult::LimitReached => println!("🔒 Turn limit reached"), EngineRunExit::Yielded => println!("↩️ Task yielded"),
EngineResult::Yielded => println!("↩️ Task yielded"), EngineRunExit::Interrupted(StopReason::LimitReached) => {
}, println!("🔒 Turn limit reached")
Err(e) => {
println!("❌ Task error: {}", e);
} }
EngineRunExit::Interrupted(reason) => println!("❌ Task interrupted: {reason:?}"), EngineRunExit::Interrupted(reason) => println!("❌ Task interrupted: {reason:?}"),
} }
crates/agen/examples/engine_cli.rs
+11 -21
View File
@@ -39,7 +39,7 @@ use tracing::info;
use tracing_subscriber::EnvFilter; use tracing_subscriber::EnvFilter;
use agen::{ use agen::{
Engine, History, Engine, EngineRunExit, StopReason,
interceptor::{Interceptor, PostToolAction, ToolResultInfo}, interceptor::{Interceptor, PostToolAction, ToolResultInfo},
llm_client::{ llm_client::{
LlmClient, LlmClient,
@@ -451,6 +451,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
// Create Engine // Create Engine
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history = agen::History::new();
let tool_call_names = Arc::new(Mutex::new(HashMap::new())); let tool_call_names = Arc::new(Mutex::new(HashMap::new()));
@@ -474,16 +475,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(&mut history, &prompt).await { let output = engine.run(&mut history, &prompt).await;
Ok(_) => {} if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = output.result {
Err(e) => { eprintln!("\n❌ Error: {error}");
eprintln!("\n❌ Error: {}", e);
std::process::exit(1);
}
} }
return Ok(()); return Ok(());
@@ -502,13 +498,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
return Ok(()); return Ok(());
} }
let mut locked = match engine.run(&mut history, first_input).await { let output = engine.run(&mut history, first_input).await;
Ok(out) => out.engine, let mut locked = output.engine;
Err(e) => {
eprintln!("\n❌ Error: {}", e);
return Ok(());
}
};
loop { loop {
print!("\n👤 You: "); print!("\n👤 You: ");
@@ -527,11 +518,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
break; break;
} }
match locked.run(&mut history, input).await { if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) =
Ok(_) => {} locked.run(&mut history, input).await
Err(e) => { {
eprintln!("\n❌ Error: {}", e); eprintln!("\n❌ Error: {error}");
}
} }
} }
crates/agen/tests/annotated_history_test.rs
+2 -3
View File
@@ -38,10 +38,9 @@ async fn run_preserves_item_annotations_without_projecting_them() {
let output = engine let output = engine
.run_with_annotation(&mut history, "hello", &mut annotate) .run_with_annotation(&mut history, "hello", &mut annotate)
.await .await;
.unwrap();
assert!(matches!(output.result, agen::EngineResult::Finished)); assert!(matches!(output.result, agen::EngineRunExit::Finished));
assert_eq!(history.len(), 2); assert_eq!(history.len(), 2);
assert_eq!(history.entries()[0].annotation, "1:user"); assert_eq!(history.entries()[0].annotation, "1:user");
assert_eq!(history.entries()[1].annotation, "2:assistant"); assert_eq!(history.entries()[1].annotation, "2:assistant");
crates/agen/tests/callback_test.rs
+21 -15
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,7 +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 mut history = agen::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();
@@ -67,7 +67,10 @@ async fn test_callback_llm_retry_event() {
}); });
let result = engine.run(&mut history, "retry once").await; let result = engine.run(&mut history, "retry once").await;
assert!(result.is_ok(), "engine should succeed after one retry"); assert!(
matches!(result.result, agen::EngineRunExit::Finished),
"engine should succeed after one retry"
);
let notices = notices.lock().unwrap(); let notices = notices.lock().unwrap();
assert_eq!(notices.len(), 1); assert_eq!(notices.len(), 1);
@@ -92,7 +95,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 mut history = agen::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()));
@@ -110,9 +113,12 @@ async fn test_callback_text_block_events() {
}); });
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineRunExit)
let result = engine.run(&mut history, "Greet me").await; let result = engine.run(&mut history, "Greet me").await;
assert!(result.is_ok(), "Engine should complete"); assert!(
matches!(result.result, agen::EngineRunExit::Finished),
"Engine should complete"
);
let deltas = text_deltas.lock().unwrap(); let deltas = text_deltas.lock().unwrap();
assert_eq!(deltas.len(), 2); assert_eq!(deltas.len(), 2);
@@ -139,7 +145,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 mut history = agen::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()));
@@ -157,7 +163,7 @@ async fn test_callback_tool_call_complete() {
}); });
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineRunExit)
let _ = engine.run(&mut history, "Weather please").await; let _ = engine.run(&mut history, "Weather please").await;
let starts = tool_starts.lock().unwrap(); let starts = tool_starts.lock().unwrap();
@@ -186,7 +192,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 mut history = agen::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()));
@@ -201,9 +207,9 @@ async fn test_callback_turn_events() {
ends.lock().unwrap().push(turn); ends.lock().unwrap().push(turn);
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineRunExit)
let result = engine.run(&mut history, "Do something").await; let result = engine.run(&mut history, "Do something").await;
assert!(result.is_ok()); assert!(matches!(result.result, agen::EngineRunExit::Finished));
let starts = turn_starts.lock().unwrap(); let starts = turn_starts.lock().unwrap();
let ends = turn_ends.lock().unwrap(); let ends = turn_ends.lock().unwrap();
@@ -258,7 +264,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(); let mut history = agen::History::new();
engine.register_tool(fixed_tool( engine.register_tool(fixed_tool(
"fixed", "fixed",
@@ -335,7 +341,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(); let mut history = agen::History::new();
engine.register_tool(erroring_tool("erroring", "boom")); engine.register_tool(erroring_tool("erroring", "boom"));
@@ -380,7 +386,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 mut history = agen::History::new();
let usage_events = Arc::new(Mutex::new(Vec::new())); let usage_events = Arc::new(Mutex::new(Vec::new()));
@@ -389,7 +395,7 @@ async fn test_callback_usage_events() {
usages.lock().unwrap().push(event.clone()); usages.lock().unwrap().push(event.clone());
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineRunExit)
let _ = engine.run(&mut history, "Hello").await; let _ = engine.run(&mut history, "Hello").await;
let usages = usage_events.lock().unwrap(); let usages = usage_events.lock().unwrap();
crates/agen/tests/engine_fixtures.rs
+4 -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,7 +134,7 @@ 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(); let mut history = agen::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(&mut history, "Hello").await; let result = engine.run(&mut history, "Hello").await;
@@ -160,7 +160,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(); let mut history = agen::History::new();
// Register tool // Register tool
let weather_tool = MockWeatherTool::new(); let weather_tool = MockWeatherTool::new();
@@ -202,7 +202,7 @@ 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(); let mut history = agen::History::new();
// Mutable::run consumes self, returns tuple // Mutable::run consumes self, returns tuple
let result = engine.run(&mut history, "Greet me").await; let result = engine.run(&mut history, "Greet me").await;
crates/agen/tests/engine_state_test.rs
+68 -100
View File
@@ -12,45 +12,12 @@ use agen::Item;
use agen::interceptor::{ use agen::interceptor::{
Interceptor, PreRequestAction, PreToolAction, ToolCallInfo, TurnEndAction, Interceptor, PreRequestAction, PreToolAction, ToolCallInfo, TurnEndAction,
}; };
use agen::llm_client::ClientError;
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, History}; use agen::{Engine, EngineError, EngineRunExit, History, StopReason};
use async_trait::async_trait; use async_trait::async_trait;
use common::MockLlmClient; use common::MockLlmClient;
#[test]
fn engine_source_has_no_worker_owned_interruption_marker() {
let source = include_str!("../src/engine.rs");
assert!(!source.contains("last_run_interrupted"));
assert!(!source.contains("set_last_run_interrupted"));
}
#[test]
fn run_exit_classifies_known_and_unexpected_stops_without_message_parsing() {
assert!(matches!(
EngineRunExit::from(Err(EngineError::Cancelled)),
EngineRunExit::Interrupted(StopReason::Cancelled)
));
assert!(matches!(
EngineRunExit::from(Err(EngineError::Client(ClientError::ContextWindowExceeded))),
EngineRunExit::Interrupted(StopReason::ContextWindowExceeded)
));
let message_only = EngineError::Client(ClientError::Api {
status: Some(400),
code: None,
message: "context_length_exceeded".to_string(),
retry_after: None,
});
assert!(matches!(
EngineRunExit::from(Err(message_only)),
EngineRunExit::Interrupted(StopReason::Unexpected(EngineError::Client(
ClientError::Api { .. }
)))
));
}
// ============================================================================= // =============================================================================
// Mutable State Tests // Mutable State Tests
// ============================================================================= // =============================================================================
@@ -235,13 +202,11 @@ async fn history_append_failure_stops_before_tool_execution() {
}); });
let mut engine = engine.lock(&history); let mut engine = engine.lock(&history);
let error = engine.run(&mut history, "use the tool").await.unwrap_err(); let exit = engine.run(&mut history, "use the tool").await;
assert!(matches!( assert!(
exit, matches!(exit, EngineRunExit::Interrupted(StopReason::Unexpected(EngineError::HistoryAppend(ref message))) if message == "simulated ENOSPC")
EngineRunExit::Interrupted(StopReason::Unexpected(EngineError::HistoryAppend(ref message))) );
if message == "simulated ENOSPC"
));
assert_eq!(tool.call_count(), 0); assert_eq!(tool.call_count(), 0);
assert_eq!(history.len(), 1); assert_eq!(history.len(), 1);
assert_eq!(history.entries()[0].item.as_text(), Some("use the tool")); assert_eq!(history.entries()[0].item.as_text(), Some("use the tool"));
@@ -319,7 +284,7 @@ async fn test_mutable_run_updates_history() -> Result<(), EngineError> {
let mut history: History = History::new(); let mut history: History = History::new();
// Execute (Mutable::run consumes self, returns EngineRunOutput) // Execute (Mutable::run consumes self, returns EngineRunOutput)
let _out = engine.run(&mut history, "Hi there").await?; let _out = engine.run(&mut history, "Hi there").await;
// History is updated // History is updated
let entries = history.entries(); let entries = history.entries();
@@ -368,12 +333,12 @@ async fn test_locked_multi_turn_history_accumulation() {
// Turn 1 // Turn 1
let result1 = locked_engine.run(&mut history, "Hello!").await; let result1 = locked_engine.run(&mut history, "Hello!").await;
assert!(result1.is_ok()); assert!(matches!(result1, EngineRunExit::Finished));
assert_eq!(history.len(), 2); // user + assistant assert_eq!(history.len(), 2); // user + assistant
// Turn 2 // Turn 2
let result2 = locked_engine.run(&mut history, "Can you help me?").await; let result2 = locked_engine.run(&mut history, "Can you help me?").await;
assert!(result2.is_ok()); assert!(matches!(result2, EngineRunExit::Finished));
assert_eq!(history.len(), 4); // 2 * (user + assistant) assert_eq!(history.len(), 4); // 2 * (user + assistant)
// Verify history contents // Verify history contents
@@ -438,10 +403,7 @@ async fn test_locked_prefix_len_tracking() {
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 locked_engine.run(&mut history, "New message").await;
.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!(history.len(), 4); // 2 + 2 assert_eq!(history.len(), 4); // 2 + 2
@@ -477,13 +439,16 @@ async fn test_turn_count_increment() -> Result<(), EngineError> {
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(&mut history, "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(&mut history, "Second").await?; assert!(matches!(
engine.run(&mut history, "Second").await,
EngineRunExit::Finished
));
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);
@@ -573,7 +538,10 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
engine.register_tool(tool_a.definition()); engine.register_tool(tool_a.definition());
let mut locked = engine.lock(&history); let mut locked = engine.lock(&history);
locked.run(&mut history, "first").await.expect("first run"); assert!(matches!(
locked.run(&mut history, "first").await,
EngineRunExit::Finished
));
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();
@@ -581,10 +549,10 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
unlocked.register_tool(tool_b.definition()); unlocked.register_tool(tool_b.definition());
let mut relocked = unlocked.lock(&history); let mut relocked = unlocked.lock(&history);
relocked assert!(matches!(
.run(&mut history, "second") relocked.run(&mut history, "second").await,
.await EngineRunExit::Finished
.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");
@@ -686,54 +654,55 @@ impl Interceptor for ContinueTurnOnce {
#[tokio::test] #[tokio::test]
async fn max_turns_is_scoped_to_each_fresh_run() { async fn max_turns_is_scoped_to_each_fresh_run() {
let mut history: History = History::new();
let responses = vec![completed_text_events(), completed_text_events()]; let 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(&history); let mut engine = engine.lock(&history);
assert_eq!( assert!(matches!(
engine.run(&mut history, "first").await.unwrap(), engine.run(&mut history, "first").await,
EngineResult::Finished EngineRunExit::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!( assert!(matches!(
engine.run(&mut history, "second").await.unwrap(), engine.run(&mut history, "second").await,
EngineResult::Finished EngineRunExit::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);
} }
#[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 history: History = History::new(); let mut history: History = History::new();
let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
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(&history); let mut engine = engine.lock(&history);
assert_eq!( assert!(matches!(
engine.run(&mut history, "start").await.unwrap(), engine.run(&mut history, "start").await,
EngineResult::Yielded EngineRunExit::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!( assert!(matches!(
engine.resume(&mut history).await.unwrap(), engine.resume(&mut history).await,
EngineResult::Finished EngineRunExit::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);
} }
#[tokio::test] #[tokio::test]
async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() { async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
let mut history: History = History::new();
let events = vec![ let events = vec![
Event::tool_use_start(0, "call_1", "count_tool"), Event::tool_use_start(0, "call_1", "count_tool"),
Event::tool_input_delta(0, "{}"), Event::tool_input_delta(0, "{}"),
@@ -744,7 +713,6 @@ 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 {
@@ -752,18 +720,18 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
}); });
let mut engine = engine.lock(&history); let mut engine = engine.lock(&history);
assert_eq!( assert!(matches!(
engine.run(&mut history, "call it").await.unwrap(), engine.run(&mut history, "call it").await,
EngineResult::Paused EngineRunExit::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!( assert!(matches!(
engine.resume(&mut history).await.unwrap(), engine.resume(&mut history).await,
EngineResult::LimitReached EngineRunExit::Interrupted(StopReason::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");
@@ -771,6 +739,7 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
#[tokio::test] #[tokio::test]
async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() { async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() {
let mut history: History = History::new();
let tool_events = vec![ let tool_events = vec![
Event::tool_use_start(0, "call_1", "count_tool"), Event::tool_use_start(0, "call_1", "count_tool"),
Event::tool_input_delta(0, "{}"), Event::tool_input_delta(0, "{}"),
@@ -782,7 +751,6 @@ 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 {
@@ -790,16 +758,16 @@ async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() {
}); });
let mut engine = engine.lock(&history); let mut engine = engine.lock(&history);
assert_eq!( assert!(matches!(
engine.run(&mut history, "pause").await.unwrap(), engine.run(&mut history, "pause").await,
EngineResult::Paused EngineRunExit::Paused
); ));
assert_eq!(engine.active_run_turn_count(), Some(1)); assert_eq!(engine.active_run_turn_count(), Some(1));
assert_eq!( assert!(matches!(
engine.run(&mut history, "replace").await.unwrap(), engine.run(&mut history, "replace").await,
EngineResult::Finished EngineRunExit::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");
@@ -807,18 +775,18 @@ 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 history: History = History::new(); let mut history: History = History::new();
let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
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(&history); let mut engine = engine.lock(&history);
assert_eq!( assert!(matches!(
engine.run(&mut history, "start").await.unwrap(), engine.run(&mut history, "start").await,
EngineResult::LimitReached EngineRunExit::Interrupted(StopReason::LimitReached)
); ));
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.llm_call_count(), 1); assert_eq!(engine.llm_call_count(), 1);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
@@ -826,17 +794,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 history: History = History::new(); let mut history: History = History::new();
let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
engine.set_turn_count(7); engine.set_turn_count(7);
engine.set_active_run_turn_count(Some(1)); engine.set_active_run_turn_count(Some(1));
let mut engine = engine.lock(&history); let mut engine = engine.lock(&history);
assert_eq!( assert!(matches!(
engine.resume(&mut history).await.unwrap(), engine.resume(&mut history).await,
EngineResult::LimitReached EngineRunExit::Interrupted(StopReason::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
+1 -4
View File
@@ -570,10 +570,7 @@ async fn test_before_tool_call_synthetic_result_committed() {
engine.set_interceptor(SyntheticPolicy); engine.set_interceptor(SyntheticPolicy);
let _result = engine let _result = engine.run(&mut history, "Test synthetic result").await;
.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!(history.items().any(|item| matches!( assert!(history.items().any(|item| matches!(
crates/agen/tests/reasoning_round_trip_test.rs
+4 -4
View File
@@ -66,7 +66,7 @@ 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 mut history: History = History::new(); let mut history: History = History::new();
let _out = engine.run(&mut history, "question?").await.expect("run ok"); let _out = engine.run(&mut history, "question?").await;
let entries = history.entries(); let entries = history.entries();
// user / reasoning / assistant_message // user / reasoning / assistant_message
@@ -110,7 +110,7 @@ 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 mut history: History = History::new(); let mut history: History = History::new();
let _out = engine.run(&mut history, "q").await.expect("run ok"); let _out = engine.run(&mut history, "q").await;
let entries = history.entries(); let entries = history.entries();
match &entries[1].item { match &entries[1].item {
@@ -156,7 +156,7 @@ 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 mut history: History = History::new(); let mut history: History = History::new();
let _out = engine.run(&mut history, "q").await.expect("run ok"); let _out = engine.run(&mut history, "q").await;
let entries = history.entries(); let entries = history.entries();
// user / reasoning(先頭) / assistant_message // user / reasoning(先頭) / assistant_message
@@ -218,7 +218,7 @@ async fn injected_reasoning_survives_into_outgoing_request() {
], ],
); );
let _ = engine.run(&mut history, "follow up").await.expect("run ok"); let _ = engine.run(&mut history, "follow up").await;
let req = captured let req = captured
.lock() .lock()
crates/session-store/tests/session_test.rs
+1 -1
View File
@@ -135,7 +135,7 @@ async fn run_and_persist(
session_id: session_store::SessionId, session_id: session_store::SessionId,
segment_id: session_store::SegmentId, segment_id: session_store::SegmentId,
input: &str, input: &str,
) -> (TestWorker, agen::EngineResult) { ) -> (TestWorker, agen::EngineRunExit) {
// 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.
crates/worker/src/worker.rs
+21 -30
View File
@@ -2941,10 +2941,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// short. Called from `Worker::run` when the worker's /// short. Called from `Worker::run` when the worker's
/// `last_run_interrupted` flag is set (i.e. the Worker just transitioned /// `last_run_interrupted` flag is set (i.e. the Worker just transitioned
/// out of Paused via a new user input). /// out of Paused via a new user input).
fn apply_interrupt_prep(&mut self) -> Result<(), WorkerError> fn apply_interrupt_prep(&mut self) -> Result<(), WorkerError> {
where
St: Clone + 'static,
{
let tool_result_summary = self let tool_result_summary = self
.prompts() .prompts()
.load_full() .load_full()
@@ -2962,17 +2959,24 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
&tool_result_summary, &tool_result_summary,
); );
if !closures.is_empty() { if !closures.is_empty() {
let mut annotate = history_annotator( let subject = worker_subject(self.session.session_id());
self.log_writer_handle(), for item in closures {
Vec::new(), let entry = HistoryEntry::new(
self.pending_committed_history.clone(), item,
); new_history_metadata(
let (engine, session) = ( WorkerHistoryProvenance::ToolOutput {
self.engine.as_mut().expect("worker present"), worker: subject.clone(),
&mut self.session, },
); None,
engine.append_history_with(session.history_mut(), closures, &mut annotate)?; ),
session.note_mutation(); );
self.commit_entry(LogEntry::AnnotatedToolResult {
ts: segment_log::now_millis(),
entry: to_logged_history_entry(&entry),
})?;
self.session.history_mut().push_entry(entry);
self.session.note_mutation();
}
} }
let interrupt_prompt_provenance = let interrupt_prompt_provenance =
self.prompt_render_provenance("internal.interrupt_system_note"); self.prompt_render_provenance("internal.interrupt_system_note");
@@ -2992,21 +2996,8 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
})?; })?;
let interrupt_entry = let interrupt_entry =
HistoryEntry::new(agen::Item::system_message(system_note), interrupt_metadata); HistoryEntry::new(agen::Item::system_message(system_note), interrupt_metadata);
let mut annotate = history_annotator( self.session.history_mut().push_entry(interrupt_entry);
self.log_writer_handle(), self.session.note_mutation();
vec![interrupt_entry.clone()],
self.pending_committed_history.clone(),
);
let (engine, session) = (
self.engine.as_mut().expect("worker present"),
&mut self.session,
);
engine.append_history_with(
session.history_mut(),
std::iter::once(interrupt_entry.item),
&mut annotate,
)?;
session.note_mutation();
Ok(()) Ok(())
} }