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/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)