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()
crates/protocol/src/lib.rs
+1
View File
@@ -923,6 +923,7 @@ pub enum WorkerStatus {
Idle,
Running,
Paused,
Stopped,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
crates/tui/src/console/mod.rs
+1 -1
View File
@@ -1016,7 +1016,7 @@ fn handle_key(app: &mut App, key: KeyEvent) -> Option<Method> {
app.clear_queued_inputs();
Some(Method::Cancel)
}
WorkerStatus::Idle => Some(Method::Shutdown),
WorkerStatus::Idle | WorkerStatus::Stopped => Some(Method::Shutdown),
}),
KeyCode::Char('d') if ctrl => {
app.quit = true;
crates/tui/src/dashboard/mod.rs
+1
View File
@@ -5204,6 +5204,7 @@ fn row_status_label(entry: &WorkerListEntry) -> (&'static str, Style) {
.fg(Color::Cyan)
.add_modifier(Modifier::BOLD),
),
Some(WorkerStatus::Stopped) => ("live stopped", Style::default().fg(Color::DarkGray)),
None => ("live", Style::default().fg(Color::DarkGray)),
};
}
crates/tui/src/workspace_panel.rs
+1
View File
@@ -1530,6 +1530,7 @@ fn worker_status_label(entry: &WorkerListEntry) -> &'static str {
Some(WorkerStatus::Idle) => "live idle",
Some(WorkerStatus::Running) => "live running",
Some(WorkerStatus::Paused) => "live paused",
Some(WorkerStatus::Stopped) => "live stopped",
None => "live",
};
}
crates/worker-runtime/src/runtime.rs
+2 -1
View File
@@ -2742,6 +2742,7 @@ impl RuntimeState {
protocol::WorkerStatus::Running => Some(WorkerStatus::Running),
protocol::WorkerStatus::Idle => Some(WorkerStatus::Idle),
protocol::WorkerStatus::Paused => Some(WorkerStatus::Paused),
protocol::WorkerStatus::Stopped => Some(WorkerStatus::Stopped),
},
protocol::Event::RunEnd { result } => match result {
protocol::RunResult::Finished | protocol::RunResult::RolledBack => {
@@ -3104,7 +3105,7 @@ mod tests {
&mut activity,
&internal_worker_status_event(
internal_worker_ref("child-b", None),
protocol::WorkerStatus::Idle,
protocol::WorkerStatus::Stopped,
),
));
}
crates/worker-runtime/src/worker_backend.rs
+3 -1
View File
@@ -1546,7 +1546,9 @@ fn accepted_notify_run_state(status: WorkerStatus, auto_run: bool) -> WorkerExec
match status {
WorkerStatus::Running => WorkerExecutionRunState::Busy,
WorkerStatus::Idle if auto_run => WorkerExecutionRunState::Busy,
WorkerStatus::Idle | WorkerStatus::Paused => WorkerExecutionRunState::Idle,
WorkerStatus::Idle | WorkerStatus::Paused | WorkerStatus::Stopped => {
WorkerExecutionRunState::Idle
}
}
}
crates/worker/src/controller.rs
+24 -4
View File
@@ -1354,7 +1354,7 @@ async fn controller_loop<C, St>(
});
}
},
WorkerStatus::Idle => {
WorkerStatus::Idle | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::NotRunning,
message: "Worker is not running".into(),
@@ -1395,7 +1395,7 @@ async fn controller_loop<C, St>(
.into(),
});
}
WorkerStatus::Running => {
WorkerStatus::Running | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message:
@@ -1409,7 +1409,7 @@ async fn controller_loop<C, St>(
WorkerStatus::Idle | WorkerStatus::Paused => {
emit_rewind_targets(&worker, &event_tx)
}
WorkerStatus::Running => {
WorkerStatus::Running | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn; rewind can only run while idle or paused"
@@ -1438,7 +1438,7 @@ async fn controller_loop<C, St>(
.into(),
});
}
WorkerStatus::Running => {
WorkerStatus::Running | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn; rewind can only run while idle or paused"
@@ -1650,6 +1650,26 @@ where
WorkerRunResult::Paused => (WorkerStatus::Paused, RunResult::Paused),
WorkerRunResult::LimitReached => (WorkerStatus::Idle, RunResult::LimitReached),
WorkerRunResult::RolledBack => (WorkerStatus::Idle, RunResult::RolledBack),
WorkerRunResult::Interrupted(_message) if pause_requested => {
let _ = event_tx.send(Event::RunEnd { result: RunResult::Paused });
return (WorkerStatus::Paused, shutdown_requested);
}
WorkerRunResult::Interrupted(message) => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: message.clone(),
});
if parent_originated {
crate::ipc::event::fire_and_forget(
parent_socket.cloned(),
protocol::WorkerEvent::Errored {
worker_name: self_name.to_string(),
message,
},
);
}
return (WorkerStatus::Idle, shutdown_requested);
}
};
let _ = event_tx.send(Event::RunEnd { result: run_result });
if parent_originated && matches!(run_result, RunResult::Finished) {
crates/worker/src/internal_worker.rs
+79 -8
View File
@@ -10,7 +10,7 @@ use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use agen::timeline::event::UsageEvent;
use agen::{Engine, llm_client::LlmClient};
use agen::{Engine, EngineError, llm_client::LlmClient};
use manifest::{Scope, WorkerManifest};
use protocol::{Event, InFlightSnapshot, WorkerStatus};
use session_store::{LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntry};
@@ -199,6 +199,20 @@ where
on_cancel_sender(worker.engine_mut().cancel_sender());
match worker.run_text(&input).await {
Ok(WorkerRunResult::Interrupted(message)) => Err(InternalWorkerError {
source: WorkerError::Engine(EngineError::Aborted(message)),
usage: last_usage.lock().ok().and_then(|slot| slot.clone()),
identity,
history_entries: store.entries_count(session_id, segment_id),
}),
Ok(WorkerRunResult::LimitReached) => Err(InternalWorkerError {
source: WorkerError::Engine(EngineError::Aborted(
"internal Worker reached its turn limit".to_string(),
)),
usage: last_usage.lock().ok().and_then(|slot| slot.clone()),
identity,
history_entries: store.entries_count(session_id, segment_id),
}),
Ok(lifecycle) => Ok(InternalWorkerResult {
usage: last_usage.lock().ok().and_then(|slot| slot.clone()),
identity,
@@ -356,7 +370,7 @@ impl InternalWorkerSessionHandle {
InternalWorkerSessionStatus::Idle => WorkerStatus::Idle,
InternalWorkerSessionStatus::Stopping
| InternalWorkerSessionStatus::Stopped
| InternalWorkerSessionStatus::Failed => WorkerStatus::Paused,
| InternalWorkerSessionStatus::Failed => WorkerStatus::Stopped,
},
error: self.last_error.lock().unwrap().clone(),
in_flight,
@@ -734,6 +748,14 @@ pub(crate) async fn prepare_internal_worker_session(
tokio::select! {
result = &mut run => {
let (turn_status, error) = match result {
Ok(WorkerRunResult::Interrupted(message)) => (
InternalWorkerSessionStatus::Stopped,
Some(message),
),
Ok(WorkerRunResult::LimitReached) => (
InternalWorkerSessionStatus::Stopped,
Some("internal Worker reached its turn limit".to_string()),
),
Ok(_) => (InternalWorkerSessionStatus::Idle, None),
Err(error) => (
InternalWorkerSessionStatus::Failed,
@@ -748,11 +770,15 @@ pub(crate) async fn prepare_internal_worker_session(
code: protocol::ErrorCode::Internal,
message,
});
} else {
let _ = event_tx.send(Event::Status {
status: WorkerStatus::Idle,
});
}
let protocol_status = if turn_status == InternalWorkerSessionStatus::Idle {
WorkerStatus::Idle
} else {
WorkerStatus::Stopped
};
let _ = event_tx.send(Event::Status {
status: protocol_status,
});
if let Some(callback) = &on_turn_end {
callback(turn_status);
}
@@ -766,7 +792,7 @@ pub(crate) async fn prepare_internal_worker_session(
let _ = (&mut run).await;
actor_in_flight.clear();
status.store(InternalWorkerSessionStatus::Stopped.encode(), std::sync::atomic::Ordering::Release);
let _ = event_tx.send(Event::Status { status: WorkerStatus::Paused });
let _ = event_tx.send(Event::Status { status: WorkerStatus::Stopped });
let _ = event_tx.send(Event::Shutdown);
state_changed.notify_waiters();
let _ = done.send(());
@@ -792,7 +818,7 @@ pub(crate) async fn prepare_internal_worker_session(
std::sync::atomic::Ordering::Release,
);
let _ = event_tx.send(Event::Status {
status: WorkerStatus::Paused,
status: WorkerStatus::Stopped,
});
let _ = event_tx.send(Event::Shutdown);
state_changed.notify_waiters();
@@ -1102,6 +1128,26 @@ mod tests {
}
}
#[derive(Clone)]
struct FailingClient;
#[async_trait]
impl LlmClient for FailingClient {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(
&self,
_request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
Err(ClientError::Config(
"intentional internal failure".to_string(),
))
}
}
#[derive(Clone)]
struct CancelBeforeAiClient {
calls: Arc<AtomicUsize>,
@@ -1215,6 +1261,31 @@ permission = "write"
assert_eq!(result.identity.kind, "test");
}
#[tokio::test]
async fn fatal_internal_run_transitions_to_stopped_protocol_status() {
let calls = Arc::new(AtomicUsize::new(0));
let mut internal_spec = spec(calls, &[]);
internal_spec.client = Box::new(FailingClient);
let handle = spawn_internal_worker_session(internal_spec)
.await
.expect("spawn failing Internal Worker session");
assert_eq!(
handle.wait_until_idle().await,
InternalWorkerSessionStatus::Stopped
);
assert_eq!(handle.status(), InternalWorkerSessionStatus::Stopped);
assert_eq!(handle.protocol_snapshot().status, WorkerStatus::Stopped);
assert!(
handle
.last_error
.lock()
.unwrap()
.as_ref()
.is_some_and(|message| message.contains("intentional internal failure"))
);
}
#[tokio::test]
async fn session_accepts_follow_up_turns_and_stops_without_runtime_registration() {
let calls = Arc::new(AtomicUsize::new(0));
crates/worker/src/spawn/tool.rs
+7 -4
View File
@@ -499,7 +499,10 @@ impl Tool for SubWorkerSpawnTool {
InternalWorkerVisibility::ParentClient,
Some(child_registry.clone()),
Some(Arc::new(move |status| {
if status == InternalWorkerSessionStatus::Failed {
if matches!(
status,
InternalWorkerSessionStatus::Failed | InternalWorkerSessionStatus::Stopped
) {
if let Some(registry) = registry.upgrade() {
if let Err(error) = registry.reclaim_internal_scope(&child_name) {
tracing::warn!(
@@ -1282,16 +1285,16 @@ extract_threshold = 4000
.unwrap();
assert_eq!(
record.session.wait_until_idle().await,
InternalWorkerSessionStatus::Failed
InternalWorkerSessionStatus::Stopped
);
assert_eq!(calls.load(Ordering::SeqCst), 3);
assert!(
spawner_scope.snapshot().is_writable(&workspace_root),
"Failed terminal child must release its delegated Workdir session"
"Stopped terminal child must release its delegated Workdir session"
);
assert!(
!record.workdir_delegation.is_active(),
"failed child must revoke cloned scoped sessions"
"stopped child must revoke cloned scoped sessions"
);
assert!(registry.get_internal("reviewer-child").is_some());
crates/worker/src/worker.rs
+92 -30
View File
@@ -10,7 +10,9 @@ use agen::llm_client::RequestConfig;
use agen::llm_client::client::LlmClient;
use agen::llm_client::types::Role;
use agen::state::Mutable;
use agen::{Engine, EngineError, EngineResult, ToolOutputLimits, UsageRecord};
use agen::{
Engine, EngineError, EngineResult, EngineRunExit, StopReason, ToolOutputLimits, UsageRecord,
};
use arc_swap::ArcSwap;
use session_store::{
LogEntry, PromptRenderProvenance, SegmentId, SessionExtension, SessionId, Store, StoreError,
@@ -905,6 +907,9 @@ pub struct Worker<C: LlmClient, St: Store> {
manifest: WorkerManifest,
/// Always `Some` outside of `run()`/`resume()`.
engine: Option<Engine<C, Mutable>>,
/// Worker-owned recovery marker. Agen exposes only the typed run exit and
/// never persists or restores Worker lifecycle state.
last_run_interrupted: bool,
store: St,
/// Optional write-through hook for name-keyed Worker metadata. Production
/// constructors install this from the same FsStore that owns the session
@@ -1107,6 +1112,7 @@ impl<C: LlmClient + Clone + 'static, St: Store + Clone + 'static> Worker<C, St>
Self {
manifest: self.manifest.clone(),
engine: Some(worker),
last_run_interrupted: false,
store: self.store.clone(),
worker_metadata_writer: None,
segment_state: self.segment_state.clone(),
@@ -1308,6 +1314,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
let mut worker = Self {
manifest,
engine: Some(worker),
last_run_interrupted: false,
store,
worker_metadata_writer: None,
segment_state: SegmentState::new(session_id, segment_id, 0),
@@ -1787,8 +1794,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
self.engine_mut().set_history(history);
self.engine_mut().set_request_config(state.config);
self.engine_mut().set_turn_count(state.turn_count);
self.engine_mut()
.set_last_run_interrupted(state.last_run_interrupted);
self.last_run_interrupted = state.last_run_interrupted;
self.engine_mut()
.set_active_run_turn_count(state.active_run_turn_count);
self.user_segments = state.user_segments;
@@ -2363,7 +2369,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
pending_attachments,
usage_history_len,
ai_activity_count: self.ai_activity_counter.load(Ordering::SeqCst),
last_run_interrupted: self.engine().last_run_interrupted(),
last_run_interrupted: self.last_run_interrupted,
active_run_turn_count: self.engine().active_run_turn_count(),
flow_runtime_state: self
.flow_runtime_state
@@ -2375,10 +2381,10 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
fn should_rollback_empty_turn(
&self,
result: &Result<EngineResult, EngineError>,
result: &EngineRunExit,
snapshot: &EmptyTurnRollbackSnapshot,
) -> bool {
if !matches!(result, Err(EngineError::Cancelled)) {
if !matches!(result, EngineRunExit::Interrupted(StopReason::Cancelled)) {
return false;
}
if self.ai_activity_counter.load(Ordering::SeqCst) != snapshot.ai_activity_count {
@@ -2394,8 +2400,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
snapshot: EmptyTurnRollbackSnapshot,
) -> Result<(), StoreError> {
self.engine_mut().truncate_history(snapshot.history_len);
self.engine_mut()
.set_last_run_interrupted(snapshot.last_run_interrupted);
self.last_run_interrupted = snapshot.last_run_interrupted;
self.engine_mut()
.set_active_run_turn_count(snapshot.active_run_turn_count);
*self
@@ -2678,9 +2683,10 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// must happen before `prepare_for_run`: proactive compaction checkpoints
/// only resumable runs, never the run this invocation is abandoning.
fn prepare_interrupted_history_for_fresh_run(&mut self) -> Result<(), WorkerError> {
if self.engine().last_run_interrupted() {
if self.last_run_interrupted {
self.apply_interrupt_prep()?;
self.engine_mut().set_last_run_interrupted(false);
self.last_run_interrupted = false;
self.engine_mut().set_active_run_turn_count(None);
}
Ok(())
}
@@ -2733,12 +2739,12 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// The explicit `PausedTurnAbandoned` marker preserves durable lifecycle
/// semantics without claiming another `run` / `resume` completed.
pub fn cancel_paused_turn(&mut self) -> Result<(), WorkerError> {
if !self.engine().last_run_interrupted() {
if !self.last_run_interrupted {
return Ok(());
}
self.apply_interrupt_prep()?;
self.engine_mut().set_last_run_interrupted(false);
self.last_run_interrupted = false;
self.commit_entry(LogEntry::PausedTurnAbandoned {
ts: segment_log::now_millis(),
})?;
@@ -2927,23 +2933,41 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// `Yielded`), so restore remains consistent.
async fn handle_worker_result(
&mut self,
result: Result<EngineResult, EngineError>,
result: EngineRunExit,
history_before: usize,
) -> Result<WorkerRunResult, WorkerError> {
self.persist_turn(history_before, &result).await?;
if matches!(result, Ok(EngineResult::Yielded)) {
if matches!(result, EngineRunExit::Yielded) {
self.last_run_interrupted = true;
return self.do_compact_and_resume().await;
}
if result.is_ok() {
if !matches!(result, EngineRunExit::Interrupted(_)) {
if let Some(ref state) = self.compact_state {
state.set_just_compacted(false);
}
}
result
.map(WorkerRunResult::from)
.map_err(WorkerError::Engine)
match result {
EngineRunExit::Finished => {
self.last_run_interrupted = false;
Ok(WorkerRunResult::Finished)
}
EngineRunExit::Paused => {
self.last_run_interrupted = true;
Ok(WorkerRunResult::Paused)
}
EngineRunExit::Interrupted(StopReason::LimitReached) => {
self.last_run_interrupted = false;
Ok(WorkerRunResult::LimitReached)
}
EngineRunExit::Interrupted(reason) => {
self.last_run_interrupted = true;
Ok(WorkerRunResult::Interrupted(stop_reason_message(&reason)))
}
EngineRunExit::Yielded => unreachable!("yielded handled above"),
}
}
fn persist_compaction_lifecycle(
@@ -3153,7 +3177,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
async fn persist_turn(
&mut self,
history_before: usize,
result: &Result<EngineResult, EngineError>,
result: &EngineRunExit,
) -> Result<(), StoreError> {
// Per-item commits for AssistantItem / ToolResult / SystemItem
// entries are expected to have landed synchronously: the
@@ -3251,22 +3275,43 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.push(record);
}
let interrupted = self.engine.as_ref().unwrap().last_run_interrupted();
let interrupted = matches!(
result,
EngineRunExit::Paused
| EngineRunExit::Yielded
| EngineRunExit::Interrupted(StopReason::Cancelled)
| EngineRunExit::Interrupted(StopReason::ContextWindowExceeded)
| EngineRunExit::Interrupted(StopReason::Unexpected(_))
);
let active_run_turn_count = self.engine.as_ref().unwrap().active_run_turn_count();
match result {
Ok(r) => {
EngineRunExit::Finished | EngineRunExit::Paused | EngineRunExit::Yielded => {
let result = match result {
EngineRunExit::Finished => EngineResult::Finished,
EngineRunExit::Paused => EngineResult::Paused,
EngineRunExit::Yielded => EngineResult::Yielded,
EngineRunExit::Interrupted(_) => unreachable!(),
};
self.commit_entry(LogEntry::RunCompleted {
ts: segment_log::now_millis(),
interrupted,
result: r.clone(),
result,
active_run_turn_count,
})?;
}
Err(e) => {
EngineRunExit::Interrupted(StopReason::LimitReached) => {
self.commit_entry(LogEntry::RunCompleted {
ts: segment_log::now_millis(),
interrupted: false,
result: EngineResult::LimitReached,
active_run_turn_count,
})?;
}
EngineRunExit::Interrupted(reason) => {
self.commit_entry(LogEntry::RunErrored {
ts: segment_log::now_millis(),
interrupted,
message: e.to_string(),
message: stop_reason_message(reason),
})?;
}
}
@@ -4467,6 +4512,9 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
fn extract_internal_worker_lifecycle_error(lifecycle: &WorkerRunResult) -> Option<WorkerError> {
match lifecycle {
WorkerRunResult::RolledBack => Some(WorkerError::Engine(EngineError::Cancelled)),
WorkerRunResult::Interrupted(message) => {
Some(WorkerError::Engine(EngineError::Aborted(message.clone())))
}
WorkerRunResult::Finished | WorkerRunResult::Paused | WorkerRunResult::LimitReached => None,
}
}
@@ -4729,6 +4777,7 @@ where
let mut worker = Self {
manifest,
engine: Some(worker),
last_run_interrupted: false,
store,
worker_metadata_writer,
segment_state: SegmentState::new(session_id, segment_id, 0),
@@ -4807,6 +4856,7 @@ where
let mut worker = Self {
manifest,
engine: Some(engine),
last_run_interrupted: false,
store,
worker_metadata_writer: None,
segment_state: SegmentState::new(session_id, segment_id, 0),
@@ -4920,6 +4970,7 @@ where
let mut worker = Self {
manifest,
engine: Some(worker),
last_run_interrupted: false,
store,
worker_metadata_writer,
segment_state: SegmentState::new(session_id, segment_id, 0),
@@ -5219,7 +5270,6 @@ where
worker.set_history(restored_history);
worker.set_request_config(state.config.clone());
worker.set_turn_count(state.turn_count);
worker.set_last_run_interrupted(state.last_run_interrupted);
worker.set_active_run_turn_count(state.active_run_turn_count);
if anchored_on_summary {
worker.set_cache_anchor(Some(0));
@@ -5234,6 +5284,7 @@ where
let mut worker = Self {
manifest,
engine: Some(worker),
last_run_interrupted: state.last_run_interrupted,
store,
worker_metadata_writer,
segment_state: SegmentState::new(session_id, segment_id, state.entries_count),
@@ -5470,8 +5521,17 @@ fn restore_manifest_from_worker_metadata_snapshot(
}
}
fn stop_reason_message(reason: &StopReason) -> String {
match reason {
StopReason::LimitReached => "engine turn limit reached".to_string(),
StopReason::ContextWindowExceeded => "model context window reached".to_string(),
StopReason::Cancelled => "engine run cancelled".to_string(),
StopReason::Unexpected(error) => format!("unexpected engine failure: {error}"),
}
}
/// Result of a Worker run.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WorkerRunResult {
/// The LLM finished its turn normally.
Finished,
@@ -5479,6 +5539,8 @@ pub enum WorkerRunResult {
Paused,
/// The worker reached its configured max_turns limit.
LimitReached,
/// The run was interrupted by a known or unexpected terminal cause.
Interrupted(String),
/// The submit-time user turn was rolled back because no AI output was materialized.
RolledBack,
}
@@ -7014,12 +7076,12 @@ mod build_summary_prompt_tests {
.await
.unwrap();
worker.ensure_segment_head().unwrap();
worker.engine_mut().set_last_run_interrupted(true);
worker.last_run_interrupted = true;
worker.engine_mut().set_active_run_turn_count(Some(3));
worker.prepare_interrupted_history_for_fresh_run().unwrap();
assert!(!worker.engine().last_run_interrupted());
assert!(!worker.last_run_interrupted);
assert_eq!(worker.engine().active_run_turn_count(), None);
let checkpoint = active_run_checkpoint_entry(
worker.engine().active_run_turn_count(),
@@ -7058,7 +7120,7 @@ mod build_summary_prompt_tests {
.unwrap();
worker.ensure_segment_head().unwrap();
worker.engine_mut().set_turn_count(7);
worker.engine_mut().set_last_run_interrupted(true);
worker.last_run_interrupted = true;
worker.engine_mut().set_active_run_turn_count(Some(3));
let session_id = worker.session_id();
@@ -7519,7 +7581,7 @@ mod build_summary_prompt_tests {
})
.unwrap();
worker.engine_mut().set_history(vec![dangling_call]);
worker.engine_mut().set_last_run_interrupted(true);
worker.last_run_interrupted = true;
worker
.run_for_notification(protocol::InvokeKind::Notify)
web/workspace/src/lib/generated/protocol.ts
+1 -1
View File
@@ -8,7 +8,7 @@ export type AlertSource = "worker" | "engine" | "compactor" | "agents_md";
export type CompletionKind = "file";
export type WorkerStatus = "idle" | "running" | "paused";
export type WorkerStatus = "idle" | "running" | "paused" | "stopped";
export type TurnResult = "finished" | "paused";