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

feat: add typed engine run exits

Browse Source
This commit is contained in:
Hare
2026-08-27 11:42:49 +09:00
parent 63306cf017
commit 975b4fa700
23 changed files with 412 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/README.md
+4 -5
View File
@@ -24,18 +24,17 @@ agen = { version = "0.2.1", features = ["codex"] }
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.
```no_run
use agen::{Engine, EngineError};
use agen::Engine;
use agen::llm_client::LlmClient;
async fn conversation<C: LlmClient>(client: C) -> Result<(), EngineError> {
async fn conversation<C: LlmClient>(client: C) {
let output = Engine::new(client)
.system_prompt("You are a concise assistant.")
.run("Explain typed state in one sentence.")
.await?;
.await;
let mut engine = output.engine;
let _result = engine.run("Give a Rust example.").await?;
Ok(())
let _exit = engine.run("Give a Rust example.").await;
}
```
crates/agen/examples/engine_cancel_demo.rs
+9 -10
View File
@@ -4,7 +4,7 @@
use agen::llm_client::scheme::{Scheme, anthropic::AnthropicScheme};
use agen::llm_client::transport::{HttpTransport, ResolvedAuth};
use agen::{Engine, EngineResult};
use agen::{Engine, EngineRunExit, StopReason};
use std::time::Duration;
#[tokio::main]
@@ -45,16 +45,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
println!("📡 Sending request to LLM...");
match engine.run("Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await {
Ok(out) => match out.result {
EngineResult::Finished => println!("✅ Task completed normally"),
EngineResult::Paused => println!("⏸️ Task paused"),
EngineResult::LimitReached => println!("🔒 Turn limit reached"),
EngineResult::Yielded => println!("↩️ Task yielded"),
},
Err(e) => {
println!("❌ Task error: {}", e);
let output = engine.run("Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await;
match output.result {
EngineRunExit::Finished => println!("✅ Task completed normally"),
EngineRunExit::Paused => println!("⏸️ Task paused"),
EngineRunExit::Yielded => println!("↩️ Task yielded"),
EngineRunExit::Interrupted(StopReason::LimitReached) => {
println!("🔒 Turn limit reached")
}
EngineRunExit::Interrupted(reason) => println!("❌ Task interrupted: {reason:?}"),
}
println!("\n✨ Demo complete!");
crates/agen/examples/engine_cli.rs
+8 -19
View File
@@ -39,7 +39,7 @@ use tracing::info;
use tracing_subscriber::EnvFilter;
use agen::{
Engine,
Engine, EngineRunExit, StopReason,
interceptor::{Interceptor, PostToolAction, ToolResultInfo},
llm_client::{
LlmClient,
@@ -476,12 +476,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
// One-shot mode
if let Some(prompt) = args.prompt {
match engine.run(&prompt).await {
Ok(_) => {}
Err(e) => {
eprintln!("\n❌ Error: {}", e);
std::process::exit(1);
}
let output = engine.run(&prompt).await;
if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = output.result {
eprintln!("\n❌ Error: {error}");
}
return Ok(());
@@ -500,13 +497,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
return Ok(());
}
let mut locked = match engine.run(first_input).await {
Ok(out) => out.engine,
Err(e) => {
eprintln!("\n❌ Error: {}", e);
return Ok(());
}
};
let output = engine.run(first_input).await;
let mut locked = output.engine;
loop {
print!("\n👤 You: ");
@@ -525,11 +517,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
break;
}
match locked.run(input).await {
Ok(_) => {}
Err(e) => {
eprintln!("\n❌ Error: {}", e);
}
if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = locked.run(input).await {
eprintln!("\n❌ Error: {error}");
}
}
crates/agen/src/engine.rs
+59 -84
View File
@@ -70,24 +70,50 @@ pub struct EngineConfig {
_private: (),
}
/// Engine execution result (status)
/// Legacy serializable outcome used by the Worker session-log compatibility boundary.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum EngineResult {
/// Completed (waiting for user input)
Finished,
/// Paused (can be resumed)
Paused,
/// Turn limit reached (max_turns exceeded)
LimitReached,
/// Yielded to caller for external processing (e.g. context compaction).
///
/// Distinct from `Paused`: internal machinery, not user-facing. The
/// caller is expected to perform some side work and then call `resume()`
/// to continue the turn loop.
Yielded,
}
/// The public termination boundary for one logical engine run.
#[derive(Debug)]
pub enum EngineRunExit {
Finished,
Paused,
Yielded,
Interrupted(StopReason),
}
/// A typed reason why an engine run could not finish normally.
#[derive(Debug)]
pub enum StopReason {
LimitReached,
ContextWindowExceeded,
Cancelled,
Unexpected(EngineError),
}
impl From<Result<EngineResult, EngineError>> for EngineRunExit {
fn from(result: Result<EngineResult, EngineError>) -> Self {
match result {
Ok(EngineResult::Finished) => Self::Finished,
Ok(EngineResult::Paused) => Self::Paused,
Ok(EngineResult::Yielded) => Self::Yielded,
Ok(EngineResult::LimitReached) => Self::Interrupted(StopReason::LimitReached),
Err(EngineError::Client(ClientError::ContextWindowExceeded)) => {
Self::Interrupted(StopReason::ContextWindowExceeded)
}
Err(EngineError::Cancelled) => Self::Interrupted(StopReason::Cancelled),
Err(error) => Self::Interrupted(StopReason::Unexpected(error)),
}
}
}
/// Result of [`Engine::run`] or [`Engine::resume`].
///
/// Contains the `Locked` Engine (ready for subsequent runs) and the outcome.
@@ -95,7 +121,7 @@ pub struct EngineRunOutput<C: LlmClient> {
/// The Engine, now in Locked state.
pub engine: Engine<C, Locked>,
/// Outcome of the turn.
pub result: EngineResult,
pub result: EngineRunExit,
}
/// Internal: tool execution result
@@ -126,16 +152,16 @@ const MAX_STREAM_CONTINUATIONS: u32 = 3;
/// engine.register_tool(my_tool);
///
/// // Mutable::run() consumes self → EngineRunOutput { engine: Locked, result }
/// let out = engine.run("Hello").await?;
/// let out = engine.run("Hello").await;
/// let mut engine = out.engine;
///
/// // Locked::run() borrows &mut self
/// engine.run("Follow-up").await?;
/// let _exit = engine.run("Follow-up").await;
///
/// // To edit between turns, unlock back to Mutable
/// let mut engine = engine.unlock();
/// engine.truncate_history(5);
/// let out = engine.run("Continue").await?;
/// let out = engine.run("Continue").await;
/// let mut engine = out.engine;
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -237,8 +263,6 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
history_append_cbs: Vec<Box<dyn Fn(&Item) -> Result<(), String> + Send + Sync>>,
/// Request configuration (max_tokens, temperature, etc.)
request_config: RequestConfig,
/// Whether the previous run was interrupted
last_run_interrupted: bool,
/// Cancel notification channel (for interrupting execution)
cancel_tx: mpsc::Sender<()>,
cancel_rx: mpsc::Receiver<()>,
@@ -270,10 +294,6 @@ pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
}
impl<C: LlmClient, S: EngineState> Engine<C, S> {
fn reset_interruption_state(&mut self) {
self.last_run_interrupted = false;
}
fn start_logical_run(&mut self) {
self.active_run_turn_count = Some(0);
}
@@ -805,11 +825,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
self.try_cancelled()
}
/// Whether the previous run was interrupted
pub fn last_run_interrupted(&self) -> bool {
self.last_run_interrupted
}
/// Generate list of ToolDefinitions for LLM from registered tools
fn build_tool_definitions(&self) -> Vec<ToolDefinition> {
self.tool_server.tool_definitions_sorted()
@@ -902,7 +917,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
match result {
Ok(value) => Ok(value),
Err(err) => {
self.last_run_interrupted = true;
let reason = match &err {
EngineError::Aborted(reason) => reason.clone(),
EngineError::Cancelled => "Cancelled".to_string(),
@@ -1000,11 +1014,9 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
continue;
}
PreToolAction::Abort(reason) => {
self.last_run_interrupted = true;
return Err(EngineError::Aborted(reason));
}
PreToolAction::Pause => {
self.last_run_interrupted = true;
return Ok(ToolExecutionResult::Paused);
}
}
@@ -1057,7 +1069,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
info!("Tool execution cancelled");
}
self.timeline.abort_current_block();
self.last_run_interrupted = true;
return Err(EngineError::Cancelled);
}
};
@@ -1079,7 +1090,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
match self.interceptor.post_tool_call(&mut info).await {
PostToolAction::Continue => {}
PostToolAction::Abort(reason) => {
self.last_run_interrupted = true;
return Err(EngineError::Aborted(reason));
}
}
@@ -1133,7 +1143,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
/// Internal turn execution logic
async fn run_turn_loop(&mut self) -> Result<EngineResult, EngineError> {
self.reset_interruption_state();
let tool_definitions = self.build_tool_definitions();
info!(
@@ -1156,7 +1165,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
if self.try_cancelled() {
info!("Execution cancelled");
self.timeline.abort_current_block();
self.last_run_interrupted = true;
return Err(EngineError::Cancelled);
}
@@ -1169,7 +1177,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
max_turns = max,
"Logical run turn limit reached"
);
self.last_run_interrupted = false;
return Ok(EngineResult::LimitReached);
}
@@ -1263,7 +1270,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
for cb in &self.turn_end_cbs {
cb(current_turn);
}
self.last_run_interrupted = true;
return Err(EngineError::Aborted(reason));
}
PreRequestAction::YieldWith(items) => {
@@ -1273,7 +1279,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
for cb in &self.turn_end_cbs {
cb(current_turn);
}
self.last_run_interrupted = true;
return Ok(EngineResult::Yielded);
}
PreRequestAction::Yield => {
@@ -1281,7 +1286,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
for cb in &self.turn_end_cbs {
cb(current_turn);
}
self.last_run_interrupted = true;
return Ok(EngineResult::Yielded);
}
PreRequestAction::ContinueWith(items) => {
@@ -1326,7 +1330,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
if let StreamCompletion::Interrupted { reason } = stream_outcome {
stream_continuations += 1;
if stream_continuations > MAX_STREAM_CONTINUATIONS {
self.last_run_interrupted = true;
return Err(EngineError::Client(ClientError::Api {
status: None,
code: None,
@@ -1378,7 +1381,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
if tool_calls.is_empty() {
match self.interceptor.on_turn_end(&self.history).await {
TurnEndAction::Finish => {
self.last_run_interrupted = false;
return Ok(EngineResult::Finished);
}
TurnEndAction::ContinueWithMessages(additional) => {
@@ -1386,7 +1388,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
continue;
}
TurnEndAction::Pause => {
self.last_run_interrupted = true;
return Ok(EngineResult::Paused);
}
}
@@ -1436,7 +1437,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
}),
);
self.timeline.abort_current_block();
self.last_run_interrupted = true;
return Err(EngineError::Cancelled);
}
};
@@ -1468,7 +1468,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
}),
);
self.timeline.abort_current_block();
self.last_run_interrupted = true;
return Err(EngineError::Cancelled);
}
};
@@ -1510,7 +1509,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
let next_failed_attempt = failed_attempt + 1;
if next_failed_attempt >= policy.max_attempts || !is_retryable(&err) {
self.last_run_interrupted = true;
return Err(EngineError::Client(err));
}
@@ -1519,7 +1517,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
.unwrap_or_else(|| policy.backoff(failed_attempt));
let elapsed = started.elapsed();
if elapsed + wait > policy.total_timeout {
self.last_run_interrupted = true;
return Err(EngineError::Client(err));
}
@@ -1548,7 +1545,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
info!("Cancelled during LLM retry backoff");
}
self.timeline.abort_current_block();
self.last_run_interrupted = true;
return Err(EngineError::Cancelled);
}
}
@@ -1591,7 +1587,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
let event = match result {
Ok(event) => event,
Err(err) => {
self.last_run_interrupted = true;
// 部分情報でも発火しておく（料金会計用）
self.timeline.flush_usage();
return Ok(StreamCompletion::Interrupted {
@@ -1612,7 +1607,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
if let Event::Error(err) = &event {
self.timeline.abort_current_block();
self.timeline.flush_usage();
self.last_run_interrupted = true;
return Err(EngineError::Client(ClientError::Api {
status: None,
code: err.code.clone(),
@@ -1630,7 +1624,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
}
self.timeline.abort_current_block();
self.timeline.flush_usage();
self.last_run_interrupted = true;
return Err(EngineError::Cancelled);
}
}
@@ -1649,10 +1642,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
tool_calls: Vec<ToolCall>,
) -> Result<Option<EngineResult>, EngineError> {
match self.execute_tools(tool_calls).await {
Ok(ToolExecutionResult::Paused) => {
self.last_run_interrupted = true;
Ok(Some(EngineResult::Paused))
}
Ok(ToolExecutionResult::Paused) => Ok(Some(EngineResult::Paused)),
Ok(ToolExecutionResult::Completed(results)) => {
// Route per-result pushes through the callback path so
// observers see each tool result as it lands.
@@ -1668,10 +1658,7 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
self.append_history_items(items)?;
Ok(None)
}
Err(err) => {
self.last_run_interrupted = true;
Err(err)
}
Err(err) => Err(err),
}
}
}
@@ -1719,7 +1706,6 @@ impl<C: LlmClient> Engine<C, Mutable> {
tool_result_cbs: Vec::new(),
history_append_cbs: Vec::new(),
request_config: RequestConfig::default(),
last_run_interrupted: false,
cancel_tx,
cancel_rx,
tool_output_limits: None,
@@ -1903,14 +1889,6 @@ impl<C: LlmClient> Engine<C, Mutable> {
self.max_turns = max_turns;
}
/// Set the last_run_interrupted flag (for session restoration)
pub fn set_last_run_interrupted(&mut self, interrupted: bool) {
self.last_run_interrupted = interrupted;
if !interrupted {
self.active_run_turn_count = None;
}
}
/// Apply configuration (reserved for future extensions)
#[allow(dead_code)]
pub fn config(self, _config: EngineConfig) -> Self {
@@ -1924,28 +1902,25 @@ impl<C: LlmClient> Engine<C, Mutable> {
///
/// Subsequent runs can call [`Engine::run`] directly.
/// To edit state between turns, call [`unlock()`](Engine::unlock) first.
pub async fn run(
self,
user_input: impl Into<String>,
) -> Result<EngineRunOutput<C>, EngineError> {
pub async fn run(self, user_input: impl Into<String>) -> EngineRunOutput<C> {
let mut locked = self.lock();
let result = locked.run(user_input).await?;
Ok(EngineRunOutput {
let result = locked.run(user_input).await;
EngineRunOutput {
engine: locked,
result,
})
}
}
/// Resume from Paused, consuming self and transitioning to Locked.
///
/// Used after `unlock()` → edit → resume.
pub async fn resume(self) -> Result<EngineRunOutput<C>, EngineError> {
pub async fn resume(self) -> EngineRunOutput<C> {
let mut locked = self.lock();
let result = locked.resume().await?;
Ok(EngineRunOutput {
let result = locked.resume().await;
EngineRunOutput {
engine: locked,
result,
})
}
}
/// Lock and transition to Locked state
@@ -1993,7 +1968,6 @@ impl<C: LlmClient> Engine<C, Mutable> {
tool_result_cbs: self.tool_result_cbs,
history_append_cbs: self.history_append_cbs,
request_config: self.request_config,
last_run_interrupted: self.last_run_interrupted,
cancel_tx: self.cancel_tx,
cancel_rx: self.cancel_rx,
@@ -2014,18 +1988,17 @@ impl<C: LlmClient> Engine<C, Locked> {
///
/// Adds a new user message to history and sends a request to the LLM.
/// Automatically loops if there are tool calls.
pub async fn run(
&mut self,
user_input: impl Into<String>,
) -> Result<EngineResult, EngineError> {
pub async fn run(&mut self, user_input: impl Into<String>) -> EngineRunExit {
self.run_result(user_input.into()).await.into()
}
async fn run_result(&mut self, user_input: String) -> Result<EngineResult, EngineError> {
// Supplying new user input abandons any paused/yielded logical run.
self.active_run_turn_count = None;
self.reset_interruption_state();
// Interceptor: on_prompt_submit
let mut user_item = Item::user_message(user_input);
let extras = match self.interceptor.on_prompt_submit(&mut user_item).await {
PromptAction::Cancel(reason) => {
self.last_run_interrupted = true;
return self
.finalize_interruption(Err(EngineError::Aborted(reason)))
.await;
@@ -2047,8 +2020,11 @@ impl<C: LlmClient> Engine<C, Locked> {
/// Resume execution (from Paused state)
///
/// Resumes turn processing from current state without adding a new user message.
pub async fn resume(&mut self) -> Result<EngineResult, EngineError> {
self.reset_interruption_state();
pub async fn resume(&mut self) -> EngineRunExit {
self.resume_result().await.into()
}
async fn resume_result(&mut self) -> Result<EngineResult, EngineError> {
self.ensure_logical_run();
let result = self.run_turn_loop().await;
let result = self.finalize_interruption(result).await;
@@ -2095,7 +2071,6 @@ impl<C: LlmClient> Engine<C, Locked> {
tool_result_cbs: self.tool_result_cbs,
history_append_cbs: self.history_append_cbs,
request_config: self.request_config,
last_run_interrupted: self.last_run_interrupted,
cancel_tx: self.cancel_tx,
cancel_rx: self.cancel_rx,
crates/agen/src/lib.rs
+2 -2
View File
@@ -20,8 +20,8 @@ pub mod usage_record;
pub use agen_macros::{description, tool, tool_registry};
pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope};
pub use engine::{
Engine, EngineConfig, EngineError, EngineResult, EngineRunOutput, LlmRetryNotice,
ToolRegistryError,
Engine, EngineConfig, EngineError, EngineResult, EngineRunExit, EngineRunOutput,
LlmRetryNotice, StopReason, ToolRegistryError,
};
pub use handler::ToolUseBlockStart;
pub use interceptor::Interceptor;
crates/agen/src/llm_client/error.rs
+8 -1
View File
@@ -18,6 +18,9 @@ pub enum ClientError {
message: String,
retry_after: Option<Duration>,
},
/// The provider rejected the request because it exceeded the model context window.
/// Classified only from a structured provider error code, never message text.
ContextWindowExceeded,
/// A request lifecycle phase exceeded its hard timeout.
Timeout {
phase: &'static str,
@@ -48,6 +51,7 @@ impl fmt::Display for ClientError {
}
write!(f, ": {}", message)
}
ClientError::ContextWindowExceeded => write!(f, "Model context window reached"),
ClientError::Timeout { phase, timeout } => {
write!(f, "{phase} timed out after {}s", timeout.as_secs())
}
@@ -112,7 +116,10 @@ pub fn is_retryable(error: &ClientError) -> bool {
ClientError::Api { status: None, .. } => false,
ClientError::Timeout { .. } => true,
ClientError::Http(e) => e.is_connect() || e.is_timeout(),
ClientError::Json(_) | ClientError::Sse(_) | ClientError::Config(_) => false,
ClientError::ContextWindowExceeded
| ClientError::Json(_)
| ClientError::Sse(_)
| ClientError::Config(_) => false,
}
}
crates/agen/src/llm_client/transport.rs
+4 -7
View File
@@ -431,13 +431,7 @@ fn api_error_code(error: &ClientError) -> Option<&str> {
}
fn is_context_length_exceeded(error: &ClientError) -> bool {
match error {
ClientError::Api { code, message, .. } => {
code.as_deref() == Some("context_length_exceeded")
|| message.contains("context_length_exceeded")
}
_ => false,
}
matches!(error, ClientError::ContextWindowExceeded)
}
async fn response_with_timeout(
@@ -487,6 +481,9 @@ async fn classify_error_response(resp: reqwest::Response) -> ClientError {
.and_then(|v| v.as_str())
.unwrap_or(&text)
.to_string();
if code.as_deref() == Some("context_length_exceeded") {
return ClientError::ContextWindowExceeded;
}
ClientError::Api {
status: Some(status),
code,
crates/agen/tests/callback_test.rs
+13 -7
View File
@@ -66,7 +66,10 @@ async fn test_callback_llm_retry_event() {
});
let result = engine.run("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();
assert_eq!(notices.len(), 1);
@@ -108,9 +111,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("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();
assert_eq!(deltas.len(), 2);
@@ -154,7 +160,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("Weather please").await;
let starts = tool_starts.lock().unwrap();
@@ -197,9 +203,9 @@ async fn test_callback_turn_events() {
ends.lock().unwrap().push(turn);
});
// Mutable::run consumes self, returns (Locked, EngineResult)
// Mutable::run consumes self, returns (Locked, EngineRunExit)
let result = engine.run("Do something").await;
assert!(result.is_ok());
assert!(matches!(result.result, agen::EngineRunExit::Finished));
let starts = turn_starts.lock().unwrap();
let ends = turn_ends.lock().unwrap();
@@ -382,7 +388,7 @@ async fn test_callback_usage_events() {
usages.lock().unwrap().push(event.clone());
});
// Mutable::run consumes self, returns (Locked, EngineResult)
// Mutable::run consumes self, returns (Locked, EngineRunExit)
let _ = engine.run("Hello").await;
let usages = usage_events.lock().unwrap();
crates/agen/tests/engine_fixtures.rs
+8 -2
View File
@@ -138,7 +138,10 @@ async fn test_engine_simple_text_response() {
// Send a simple message (Mutable::run consumes self, returns tuple)
let result = engine.run("Hello").await;
assert!(result.is_ok(), "Engine should complete successfully");
assert!(
matches!(result.result, agen::EngineRunExit::Finished),
"Engine should complete successfully"
);
}
/// Verify that Engine can correctly process responses containing tool calls
@@ -199,7 +202,10 @@ async fn test_engine_with_programmatic_events() {
// Mutable::run consumes self, returns tuple
let result = engine.run("Greet me").await;
assert!(result.is_ok(), "Engine should complete successfully");
assert!(
matches!(result.result, agen::EngineRunExit::Finished),
"Engine should complete successfully"
);
}
/// Verify that ToolCallCollector correctly collects ToolCall from ToolUse block events
crates/agen/tests/engine_state_test.rs
+73 -27
View File
@@ -12,12 +12,45 @@ use agen::Item;
use agen::interceptor::{
Interceptor, PreRequestAction, PreToolAction, ToolCallInfo, TurnEndAction,
};
use agen::llm_client::ClientError;
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, EngineError, EngineResult};
use agen::{Engine, EngineError, EngineRunExit, StopReason};
use async_trait::async_trait;
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
// =============================================================================
@@ -195,11 +228,13 @@ async fn history_append_failure_stops_before_tool_execution() {
});
let mut engine = engine.lock();
let error = engine.run("use the tool").await.unwrap_err();
let exit = engine.run("use the tool").await;
assert!(
matches!(error, EngineError::HistoryAppend(ref message) if message == "simulated ENOSPC")
);
assert!(matches!(
exit,
EngineRunExit::Interrupted(StopReason::Unexpected(EngineError::HistoryAppend(ref message)))
if message == "simulated ENOSPC"
));
assert_eq!(tool.call_count(), 0);
assert_eq!(engine.history().len(), 1);
assert_eq!(engine.history()[0].as_text(), Some("use the tool"));
@@ -274,7 +309,7 @@ async fn test_mutable_run_updates_history() -> Result<(), EngineError> {
let engine = Engine::new(client);
// Execute (Mutable::run consumes self, returns EngineRunOutput)
let out = engine.run("Hi there").await?;
let out = engine.run("Hi there").await;
let engine = out.engine;
// History is updated
@@ -323,12 +358,12 @@ async fn test_locked_multi_turn_history_accumulation() {
// Turn 1
let result1 = locked_engine.run("Hello!").await;
assert!(result1.is_ok());
assert!(matches!(result1, EngineRunExit::Finished));
assert_eq!(locked_engine.history().len(), 2); // user + assistant
// Turn 2
let result2 = locked_engine.run("Can you help me?").await;
assert!(result2.is_ok());
assert!(matches!(result2, EngineRunExit::Finished));
assert_eq!(locked_engine.history().len(), 4); // 2 * (user + assistant)
// Verify history contents
@@ -386,7 +421,7 @@ async fn test_locked_prefix_len_tracking() {
assert_eq!(locked_engine.locked_prefix_len(), 2); // 2 items at lock time
// Execute turn
locked_engine.run("New message").await.unwrap();
locked_engine.run("New message").await;
// History grows but locked_prefix_len remains unchanged
assert_eq!(locked_engine.history().len(), 4); // 2 + 2
@@ -421,13 +456,13 @@ async fn test_turn_count_increment() -> Result<(), EngineError> {
assert_eq!(engine.llm_call_count(), 0);
// First run consumes Mutable, returns EngineRunOutput
let mut engine = engine.run("First").await?.engine;
let mut engine = engine.run("First").await.engine;
assert_eq!(engine.turn_count(), 1);
// Retry not yet implemented → AgentTurn:LlmCall is 1:1.
assert_eq!(engine.llm_call_count(), 1);
// Subsequent runs on Locked take &mut self
engine.run("Second").await?;
engine.run("Second").await;
assert_eq!(engine.turn_count(), 2);
assert_eq!(engine.llm_call_count(), 2);
@@ -515,7 +550,7 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
engine.register_tool(tool_a.definition());
let mut locked = engine.lock();
locked.run("first").await.expect("first run");
locked.run("first").await;
assert_eq!(tool_a.call_count(), 1, "tool_a should be called once");
let mut unlocked = locked.unlock();
@@ -523,7 +558,7 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
unlocked.register_tool(tool_b.definition());
let mut relocked = unlocked.lock();
relocked.run("second").await.expect("second run");
relocked.run("second").await;
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");
@@ -628,11 +663,14 @@ async fn max_turns_is_scoped_to_each_fresh_run() {
engine.set_max_turns(Some(1));
let mut engine = engine.lock();
assert_eq!(engine.run("first").await.unwrap(), EngineResult::Finished);
assert!(matches!(engine.run("first").await, EngineRunExit::Finished));
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(engine.run("second").await.unwrap(), EngineResult::Finished);
assert!(matches!(
engine.run("second").await,
EngineRunExit::Finished
));
assert_eq!(engine.turn_count(), 2);
assert_eq!(engine.active_run_turn_count(), None);
}
@@ -646,11 +684,11 @@ async fn yielded_resume_keeps_the_same_unspent_turn_budget() {
});
let mut engine = engine.lock();
assert_eq!(engine.run("start").await.unwrap(), EngineResult::Yielded);
assert!(matches!(engine.run("start").await, EngineRunExit::Yielded));
assert_eq!(engine.turn_count(), 0);
assert_eq!(engine.active_run_turn_count(), Some(0));
assert_eq!(engine.resume().await.unwrap(), EngineResult::Finished);
assert!(matches!(engine.resume().await, EngineRunExit::Finished));
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None);
}
@@ -674,12 +712,15 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
});
let mut engine = engine.lock();
assert_eq!(engine.run("call it").await.unwrap(), EngineResult::Paused);
assert!(matches!(engine.run("call it").await, EngineRunExit::Paused));
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), Some(1));
assert_eq!(tool.call_count(), 0);
assert_eq!(engine.resume().await.unwrap(), EngineResult::LimitReached);
assert!(matches!(
engine.resume().await,
EngineRunExit::Interrupted(StopReason::LimitReached)
));
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(tool.call_count(), 1, "the consumed turn's tool still runs");
@@ -705,10 +746,13 @@ async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() {
});
let mut engine = engine.lock();
assert_eq!(engine.run("pause").await.unwrap(), EngineResult::Paused);
assert!(matches!(engine.run("pause").await, EngineRunExit::Paused));
assert_eq!(engine.active_run_turn_count(), Some(1));
assert_eq!(engine.run("replace").await.unwrap(), EngineResult::Finished);
assert!(matches!(
engine.run("replace").await,
EngineRunExit::Finished
));
assert_eq!(engine.turn_count(), 2);
assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(tool.call_count(), 1, "pending-tool semantics are unchanged");
@@ -723,10 +767,10 @@ async fn interceptor_continuation_consumes_the_logical_run_budget() {
});
let mut engine = engine.lock();
assert_eq!(
engine.run("start").await.unwrap(),
EngineResult::LimitReached
);
assert!(matches!(
engine.run("start").await,
EngineRunExit::Interrupted(StopReason::LimitReached)
));
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.llm_call_count(), 1);
assert_eq!(engine.active_run_turn_count(), None);
@@ -737,11 +781,13 @@ async fn restored_active_run_budget_is_enforced_before_another_llm_call() {
let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
engine.set_max_turns(Some(1));
engine.set_turn_count(7);
engine.set_last_run_interrupted(true);
engine.set_active_run_turn_count(Some(1));
let mut engine = engine.lock();
assert_eq!(engine.resume().await.unwrap(), EngineResult::LimitReached);
assert!(matches!(
engine.resume().await,
EngineRunExit::Interrupted(StopReason::LimitReached)
));
assert_eq!(engine.turn_count(), 7);
assert_eq!(engine.llm_call_count(), 0);
assert_eq!(engine.active_run_turn_count(), None);
crates/agen/tests/parallel_execution_test.rs
+8 -5
View File
@@ -158,7 +158,7 @@ async fn test_parallel_tool_execution() {
engine.register_tool(tool3.definition());
let start = Instant::now();
// Mutable::run consumes self, returns (Locked, EngineResult)
// Mutable::run consumes self, returns (Locked, EngineRunExit)
let _result = engine.run("Run all tools").await;
let elapsed = start.elapsed();
@@ -415,7 +415,7 @@ async fn test_before_tool_call_skip() {
engine.set_interceptor(BlockingPolicy);
// Mutable::run consumes self, returns (Locked, EngineResult)
// Mutable::run consumes self, returns (Locked, EngineRunExit)
let _result = engine.run("Test hook").await;
// allowed_tool is called, but blocked_tool is not
@@ -502,10 +502,13 @@ async fn test_post_tool_call_modification() {
modified_content: modified_content.clone(),
});
// Mutable::run consumes self, returns (Locked, EngineResult)
// Mutable::run consumes self, returns (Locked, EngineRunExit)
let result = engine.run("Test modification").await;
assert!(result.is_ok(), "Engine should complete");
assert!(
matches!(result.result, agen::EngineRunExit::Finished),
"Engine should complete"
);
// Verify hook was called and content was modified
let content = modified_content.lock().unwrap().clone();
@@ -558,7 +561,7 @@ async fn test_before_tool_call_synthetic_result_committed() {
engine.set_interceptor(SyntheticPolicy);
let result = engine.run("Test synthetic result").await.unwrap();
let result = engine.run("Test synthetic result").await;
assert_eq!(blocked_clone.call_count(), 0, "Blocked tool should not run");
assert!(result.engine.history().iter().any(|item| matches!(
crates/agen/tests/reasoning_round_trip_test.rs
+4 -4
View File
@@ -65,7 +65,7 @@ async fn anthropic_thinking_round_trips_signature_into_history() {
]);
let client = MockLlmClient::new(events);
let engine = Engine::new(client);
let out = engine.run("question?").await.expect("run ok");
let out = engine.run("question?").await;
let engine = out.engine;
let history = engine.history();
@@ -109,7 +109,7 @@ async fn openai_reasoning_round_trips_encrypted_and_summary() {
]);
let client = MockLlmClient::new(events);
let engine = Engine::new(client);
let out = engine.run("q").await.expect("run ok");
let out = engine.run("q").await;
let engine = out.engine;
let history = engine.history();
@@ -155,7 +155,7 @@ async fn reasoning_precedes_text_in_assistant_burst() {
}));
let client = MockLlmClient::new(events);
let engine = Engine::new(client);
let out = engine.run("q").await.expect("run ok");
let out = engine.run("q").await;
let engine = out.engine;
let history = engine.history();
@@ -214,7 +214,7 @@ async fn injected_reasoning_survives_into_outgoing_request() {
Item::assistant_message("prior answer"),
]);
let _ = engine.run("follow up").await.expect("run ok");
let _ = engine.run("follow up").await;
let req = captured
.lock()