From 22179a34822660d1bce25fd5b7ec57cf2dd93223 Mon Sep 17 00:00:00 2001 From: daniel Date: Thu, 8 Oct 2026 00:01:34 +0100 Subject: [PATCH 1/3] feat(loop): add awaited mutator lifecycle callbacks --- book/src/ch16-compaction.md | 2 +- crates/agentkit-loop/README.md | 14 + crates/agentkit-loop/src/lib.rs | 1459 +++++++++++++++++-- crates/agentkit-loop/src/lifecycle_tests.rs | 1415 ++++++++++++++++++ 4 files changed, 2764 insertions(+), 126 deletions(-) create mode 100644 crates/agentkit-loop/src/lifecycle_tests.rs diff --git a/book/src/ch16-compaction.md b/book/src/ch16-compaction.md index 913421e..bfacc54 100644 --- a/book/src/ch16-compaction.md +++ b/book/src/ch16-compaction.md @@ -261,7 +261,7 @@ This is why caching is configured separately from compaction in agentkit. Compac ## Loop integration -Compactors register as `LoopMutator`s. The loop runs every registered mutator at each `MutationPoint` — `AfterToolResult` (between tool results and the next inference call) and `AfterTurnEnded` (after the assistant final, interrupt, or cancellation). The trigger decides which points are relevant. +Compactors register as `LoopMutator`s. The loop runs transcript mutators before inference at `MutationPoint::AfterToolResult` (when continuing after a tool round) or `MutationPoint::AfterTurnEnded` (when dispatching a new logical turn). Despite its name, `AfterTurnEnded` is not a completion notification: newly buffered input has already been appended when it runs. The trigger decides which points are relevant. When a compactor fires: diff --git a/crates/agentkit-loop/README.md b/crates/agentkit-loop/README.md index 9902283..d47a981 100644 --- a/crates/agentkit-loop/README.md +++ b/crates/agentkit-loop/README.md @@ -19,6 +19,20 @@ This crate provides: Use it as the central coordinator between model providers, tool execution, and application UI or control flow. +## Awaited lifecycle extensions + +Register a `LoopMutator` with `AgentBuilder::mutator` to intercept lifecycle boundaries without driving or copying the loop implementation. Its lifecycle methods are asynchronous and default to no-ops, so existing transcript mutators remain compatible. Registered callbacks run in order and later callbacks see earlier payload edits. + +The extension surface covers session startup and explicit closure, input admission, logical-turn startup and finalization, model requests and complete responses, model errors and stream progress, and resolved top-level tool batches. A logical turn may include multiple model requests and tool continuations; turn callbacks do not run once per inference. + +Payload types separate editable content from read-only identity and accounting. Complete response edits are applied before transcript commit and tool dispatch. Final-response edits are reflected in both the transcript and returned turn result. Stream-progress callbacks do not retract deltas already delivered to observers. + +Use `LoopDriver::submit_input_async` when input must pass awaited callbacks before acceptance. Synchronous submission defers that work until the next driver step. Use `LoopDriver::close().await` to finalize an owned session: a prompt finishing is not session closure, and dropping a driver cannot run asynchronous callbacks. + +Lifecycle extensions complement existing boundaries rather than replacing them: decorate a `ToolExecutor` for individual or nested tool interception, and a compaction implementation for an actual compaction operation. Model-selection and application-owned configuration changes remain outside the loop. + +The original `mutate` callback still runs at the two pre-inference transcript mutation points. `AfterTurnEnded` names new-turn dispatch, not a completion notification; `AfterToolResult` names a tool-round continuation. Use lifecycle callbacks for actual start/end notifications. + ## Quick start ```rust,no_run diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index 33f7c43..28fe115 100644 --- a/crates/agentkit-loop/src/lib.rs +++ b/crates/agentkit-loop/src/lib.rs @@ -792,9 +792,9 @@ pub struct ObservedEvent { /// replication, or audit. /// /// Observers are called *synchronously* from inside the driver, in the -/// same order items land in the transcript. Compaction-driven transcript -/// rewrites do **not** fire `on_transcript_event` — those are signaled by -/// [`AgentEvent::CompactionFinished`] instead. +/// same order items land in the transcript. Replacements are delivered through +/// [`TranscriptObserver::on_transcript_rewrite`], not as fictitious appends. +/// Persistence consumers that enable transcript editing must handle both methods. /// /// Register via [`AgentBuilder::transcript_observer`]; multiple observers /// may be registered and are called in registration order. @@ -820,6 +820,10 @@ pub trait TranscriptObserver: Send + Sync { /// state behind interior mutability so the driver can share an /// `Arc`. fn on_transcript_event(&self, event: TranscriptEvent<'_>); + + /// Replace the persisted transcript with this complete committed snapshot. + /// The default preserves compatibility for append-only observers. + fn on_transcript_rewrite(&self, _event: TranscriptRewriteEvent<'_>) {} } /// Session-addressed transcript append event delivered to @@ -832,6 +836,13 @@ pub struct TranscriptEvent<'a> { pub item: &'a Item, } +/// A complete committed replacement, including staged turn-start input. +#[derive(Clone, Debug)] +pub struct TranscriptRewriteEvent<'a> { + pub session_id: &'a SessionId, + pub items: &'a [Item], +} + /// Where in the loop a [`LoopMutator`] is given a chance to modify the /// transcript. Mutators run synchronously at these points; mid-stream /// mutation (e.g. between content deltas) is intentionally not supported @@ -842,8 +853,9 @@ pub enum MutationPoint { /// A tool result has just been appended; the next loop step will be /// another inference call. AfterToolResult, - /// A turn has fully ended (assistant final, interrupt, or cancellation) - /// and any new user input has not yet been dispatched. + /// Legacy name for the pre-inference transcript mutation pass on a new + /// logical turn, after input append. This is not a terminal callback; + /// use [`LoopMutator::on_turn_end`] for actual logical retirement. AfterTurnEnded, } @@ -895,15 +907,233 @@ impl<'a> std::ops::DerefMut for TranscriptCursor<'a> { } } +/// Read-only identity and cancellation context for awaited lifecycle callbacks. +/// +/// A logical turn can contain several model calls and approval/tool continuations. +/// `model_call_index` is session-local and counts calls, not provider retry attempts. +/// Terminal callbacks run even when the cancellation token is cancelled. +#[non_exhaustive] +pub struct LifecycleCtx<'a> { + pub session_id: &'a SessionId, + pub turn_id: Option<&'a agentkit_core::TurnId>, + pub model_call_index: Option, + pub cancellation: Option, + /// Selection snapshot; per-call routing inside `begin_turn` can change it. + pub provider_name: Option<&'a str>, + pub model_name: Option<&'a str>, + /// Original failure at a terminal boundary, never replaced by hook errors. + pub error: Option<&'a LoopError>, +} + +/// Read-only state at explicit session closure. No detached work is awaited. +pub struct SessionEndPayload<'a> { + pub transcript: &'a [Item], +} + +/// Loop-owned logical progress. This is not arbitrary progress reported by tools. +pub enum TurnProgress<'a> { + Model(&'a ModelTurnEvent), + ToolDetached(&'a ToolResultPart), +} + +/// Editable session options; the session identity is available only in the context. +pub struct SessionStartPayload<'a> { + pub metadata: &'a mut MetadataMap, + pub cache: &'a mut Option, + pub consumer_capabilities: &'a mut SessionConsumerCapabilities, +} + +/// Inference-local request options. Transcript edits affect only this model call; +/// persistent history edits use [`LoopMutator::mutate`] or `on_turn_start`. +/// Catalog changes do not install executable tools. Session/turn identity is read-only. +/// Prompt edits are validated before the request is consumed. +pub struct ModelRequestPayload<'a> { + pub transcript: &'a mut Vec, + pub available_tools: &'a mut Vec, + pub cache: &'a mut Option, + pub metadata: &'a mut MetadataMap, +} + +/// Editable content of one output item. Identity, role, accounting and timestamp +/// remain read-only in `original`, the snapshot at entry to this callback. +/// `parts` and `metadata` expose the current candidate, including earlier hooks. +/// Tool linkage in `parts` is checked before commit. +pub struct OutputItemPayload<'a> { + pub original: &'a Item, + pub parts: &'a mut Vec, + pub metadata: &'a mut MetadataMap, +} + +/// Complete, uncommitted model output. Item identities, roles, accounting, timestamps, +/// and tool-call/result identities must remain unchanged; content can be rewritten. +/// The driver checks these constraints before committing any output. +/// Already emitted streaming deltas cannot be retracted by this callback. +pub struct ModelResponsePayload<'a> { + pub output: Vec>, + pub finish_reason: &'a FinishReason, + pub usage: Option<&'a Usage>, + pub metadata: &'a MetadataMap, + pub model: Option<&'a str>, + pub response_id: Option<&'a str>, +} + +/// Logical finalization candidate. `result` is the immutable snapshot at entry +/// to this callback, including earlier hooks' edits. +/// `output` is editable only for uncommitted successful final assistant output. +/// Cancellation, failure, and already committed tool-bearing output are read-only. +/// Output edits obey the same constraints as [`ModelResponsePayload`]. +pub struct TurnFinishPayload<'a> { + pub result: &'a TurnResult, + pub output: Option>>, +} + +/// A transcript-resolved top-level tool batch, not a nested execution batch. +/// Results include cancellation/denial and background-detachment placeholders; +/// a placeholder does not imply that background execution has completed. +pub struct ToolBatchPayload<'a> { + /// Model-issued calls. Approved argument rewrites belong to executor facts. + pub calls: &'a [ToolCallPart], + pub results: &'a [ToolResultPart], +} + /// Async transcript mutator. Registered via [`AgentBuilder::mutator`] and /// invoked at each [`MutationPoint`]. Mutators own their derived state /// (e.g. running token totals via interior mutability) and decide for /// themselves whether and how to modify the transcript. /// -/// The default implementation is a no-op so trait users override only -/// `mutate`. +/// Every method defaults to a no-op. Lifecycle callbacks are awaited in +/// registration order at their typed boundaries; `mutate` remains the generic +/// transcript editing seam. Staged edits are committed only after validation. +/// Tool/executor, compaction, routing and application configuration remain +/// separate extension seams; this trait does not intercept nested tools. #[async_trait] pub trait LoopMutator: Send + Sync { + /// Runs before session options are consumed by the model adapter. + async fn on_session_start( + &self, + _payload: SessionStartPayload<'_>, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Runs once on explicit [`LoopDriver::close`], not on turn completion or Drop. + async fn on_session_end( + &self, + _payload: SessionEndPayload<'_>, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Runs before input dispatch. Sync submission queues admission until `next`; + /// [`LoopDriver::submit_input_async`] awaits this callback before acceptance. + async fn on_input( + &self, + _input: &mut Vec, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Runs once before work in a newly active logical turn. The cursor is a + /// staged transcript including admitted input. Edits are validated before + /// commit; failure leaves the live transcript unchanged and retires the turn. + async fn on_turn_start( + &self, + _transcript: &mut TranscriptCursor<'_>, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Runs before logical finalization, including cancellation/error (read-only). + async fn on_turn_finish( + &self, + _payload: TurnFinishPayload<'_>, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Runs once after logical retirement. Every registered callback is attempted. + /// On ordinary retirement, errors are logged without replacing the committed + /// result. Explicit `close` aggregates notification errors in its close result. + async fn on_turn_end( + &self, + _result: &TurnResult, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Runs immediately before the inference request is consumed by the session. + async fn on_model_request( + &self, + _payload: ModelRequestPayload<'_>, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Runs on complete output, before transcript commit and tool dispatch. + async fn on_model_response( + &self, + _payload: ModelResponsePayload<'_>, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Read-only model failure notification, including cancellation. Hook errors + /// never replace the model error. This does not report tool or hook failures. + async fn on_model_error( + &self, + _error: &LoopError, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Read-only logical progress before forwarding model events or committing + /// a background-detachment placeholder. Individual tool progress is external. + async fn on_turn_progress( + &self, + _progress: TurnProgress<'_>, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Read-only notification of loop-owned background-detachment progress. + async fn on_tool_progress( + &self, + _progress: &ToolResultPart, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Read-only stream/progress notification, before observer forwarding. + /// Includes usage, tool requests, Finished, and response-attempt supersession. + async fn on_model_progress( + &self, + _event: &ModelTurnEvent, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + + /// Runs once when every top-level call in a model batch has a transcript result. + /// Individual and nested tool interception belongs at [`ToolExecutor`]. + async fn on_tool_batch( + &self, + _payload: ToolBatchPayload<'_>, + _ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + Ok(()) + } + /// Run this mutator. Returning without writing to `cursor` is a no-op. /// Errors abort the loop; protocol-violating mutations (orphaned tool /// uses or results) are detected by validation and turned into @@ -1313,9 +1543,47 @@ struct ActiveToolRound { foreground_progressed: bool, } +struct PendingToolBatch { + turn_id: agentkit_core::TurnId, + calls: Vec, + results: BTreeMap, + next_mutator: usize, + error: Option, +} + +struct TerminalState { + result: TurnResult, + original_items: Vec, + editable: bool, + output_committed: bool, + finish_index: usize, + finish_done: bool, + terminal_emitted: bool, + end_index: usize, + failure: Option, + finish_error: Option, + end_error: Option, +} + +struct ClosingState { + cleanup_turns: Vec, + cleanup_index: usize, + cleanup_done: bool, + session_index: usize, + error: Option, +} + #[derive(Default)] struct DriverLifecycle { active_turn: Option, + closed: bool, + closing: Option, + terminal: Option, + model_call_index: u64, + pending_cancel: Option, + finish_attempted: bool, + cancellation: Option, + started_with_input: bool, } /// A configured agent ready to start a session. @@ -1402,8 +1670,50 @@ where /// # Errors /// /// Returns [`LoopError`] if the model adapter fails to create a session. - pub async fn start(&self, config: SessionConfig) -> Result, LoopError> { + pub async fn start( + &self, + mut config: SessionConfig, + ) -> Result, LoopError> { let session_id = config.session_id.clone(); + for mutator in &self.mutators { + mutator + .on_session_start( + SessionStartPayload { + metadata: &mut config.metadata, + cache: &mut config.cache, + consumer_capabilities: &mut config.consumer_capabilities, + }, + LifecycleCtx { + session_id: &session_id, + turn_id: None, + model_call_index: None, + cancellation: None, + provider_name: self.model.provider_name(), + model_name: None, + error: None, + }, + ) + .await?; + } + let mut input = self.input.clone(); + if !input.is_empty() { + for mutator in &self.mutators { + mutator + .on_input( + &mut input, + LifecycleCtx { + session_id: &session_id, + turn_id: None, + model_call_index: None, + cancellation: None, + provider_name: self.model.provider_name(), + model_name: None, + error: None, + }, + ) + .await?; + } + } let default_cache = config.cache.clone(); let mut session = self.model.start_session(config).await?; if !self.observers.is_empty() { @@ -1439,7 +1749,9 @@ where observers: self.observers.clone(), transcript_observers: self.transcript_observers.clone(), transcript: self.transcript.clone(), - pending_input: self.input.clone(), + pending_input: input, + pending_admission: VecDeque::new(), + pending_tool_batches: Vec::new(), pending_approvals: BTreeMap::new(), pending_approval_order: VecDeque::new(), active_tool_round: None, @@ -1453,6 +1765,12 @@ where tool_cancellations: HashMap::new(), }; driver.emit(AgentEvent::RunStarted { session_id }); + if !self.mutators.is_empty() && !driver.pending_input.is_empty() { + driver.emit(AgentEvent::InputAccepted { + session_id: driver.session_id.clone(), + items: driver.pending_input.clone(), + }); + } Ok(driver) } } @@ -1715,6 +2033,8 @@ where transcript_observers: Vec>, transcript: Vec, pending_input: Vec, + pending_admission: VecDeque>, + pending_tool_batches: Vec, pending_approvals: BTreeMap, pending_approval_order: VecDeque, active_tool_round: Option, @@ -1843,34 +2163,332 @@ where !self.pending_approvals.is_empty() } - fn start_logical_turn(&mut self) -> agentkit_core::TurnId { + async fn start_logical_turn(&mut self) -> Result { if let Some(turn_id) = &self.lifecycle.active_turn { - return turn_id.clone(); + return Ok(turn_id.clone()); } let turn_id = agentkit_core::TurnId::new(format!("turn-{}", self.next_turn_index)); self.next_turn_index += 1; - self.start_logical_turn_with(turn_id) + self.start_logical_turn_with(turn_id).await } - fn start_logical_turn_with(&mut self, turn_id: agentkit_core::TurnId) -> agentkit_core::TurnId { - if let Some(active_turn) = &self.lifecycle.active_turn { - return active_turn.clone(); + async fn start_logical_turn_with( + &mut self, + turn_id: agentkit_core::TurnId, + ) -> Result { + if self.lifecycle.active_turn.as_ref() == Some(&turn_id) { + return Ok(turn_id); } self.lifecycle.active_turn = Some(turn_id.clone()); self.emit(AgentEvent::TurnStarted { session_id: self.session_id.clone(), turn_id: turn_id.clone(), }); - turn_id + self.lifecycle.finish_attempted = false; + self.lifecycle.cancellation = self + .cancellation + .as_ref() + .map(CancellationHandle::checkpoint); + if !self.mutators.is_empty() { + let mut candidate = self.transcript.clone(); + self.lifecycle.started_with_input = !self.pending_input.is_empty(); + candidate.extend(std::mem::take(&mut self.pending_input)); + let original = candidate.clone(); + let mut cursor = TranscriptCursor { + items: &mut candidate, + dirty: false, + }; + for mutator in &self.mutators { + mutator + .on_turn_start(&mut cursor, self.lifecycle_ctx(Some(&turn_id), None)) + .await?; + } + validate_prompt_edits(&original, &candidate)?; + self.commit_transcript(candidate); + } + Ok(turn_id) } - fn finish_logical_turn(&mut self, result: &TurnResult) { + async fn finish_logical_turn( + &mut self, + result: &TurnResult, + failure: Option<&LoopError>, + ) -> Result<(), LoopError> { if self.pending_round_resume.as_ref() == Some(&result.turn_id) { self.pending_round_resume = None; } - if self.lifecycle.active_turn.as_ref() == Some(&result.turn_id) { + if self.lifecycle.terminal.is_none() { + self.begin_terminal(result, false, failure, false); + self.lifecycle + .terminal + .as_mut() + .expect("terminal state") + .finish_done = true; + } + let state = self.lifecycle.terminal.as_mut().expect("terminal state"); + state.result = result.clone(); + state.output_committed = true; + if let Some(error) = failure { + state.failure = Some(error.clone()); + } + if !state.terminal_emitted { + if self.lifecycle.active_turn.as_ref() == Some(&result.turn_id) { + self.emit(AgentEvent::TurnFinished(result.clone())); + } self.lifecycle.active_turn = None; - self.emit(AgentEvent::TurnFinished(result.clone())); + self.lifecycle.started_with_input = false; + self.lifecycle + .terminal + .as_mut() + .expect("terminal state") + .terminal_emitted = true; + } + loop { + let state = self.lifecycle.terminal.as_ref().expect("terminal state"); + if state.end_index == self.mutators.len() { + break; + } + let outcome = self.mutators[state.end_index] + .on_turn_end( + &state.result, + LifecycleCtx { + error: state.failure.as_ref(), + ..self.lifecycle_ctx(Some(&state.result.turn_id), None) + }, + ) + .await; + let state = self.lifecycle.terminal.as_mut().expect("terminal state"); + state.end_index += 1; + if let Err(error) = outcome { + state.end_error.get_or_insert(error); + } + } + let state = self.lifecycle.terminal.take().expect("terminal state"); + if self.lifecycle.closing.is_some() { + return state.end_error.map_or(Ok(()), Err); + } + if let Some(error) = state.end_error { + tracing::debug!(%error, "turn end callback failed after retirement"); + } + Ok(()) + } + + fn begin_terminal( + &mut self, + result: &TurnResult, + editable: bool, + failure: Option<&LoopError>, + append_output: bool, + ) { + self.lifecycle.terminal = Some(TerminalState { + result: result.clone(), + original_items: result.items.clone(), + editable, + output_committed: !append_output, + finish_index: 0, + finish_done: false, + terminal_emitted: false, + end_index: 0, + failure: failure.cloned(), + finish_error: None, + end_error: None, + }); + } + + fn mark_terminal_output_committed(&mut self, turn_id: &agentkit_core::TurnId) { + if let Some(state) = &mut self.lifecycle.terminal { + if &state.result.turn_id == turn_id { + state.output_committed = true; + } + } + } + + fn lifecycle_ctx<'a>( + &'a self, + turn_id: Option<&'a agentkit_core::TurnId>, + model_call_index: Option, + ) -> LifecycleCtx<'a> { + LifecycleCtx { + session_id: &self.session_id, + turn_id, + model_call_index, + cancellation: if turn_id.is_some() { + self.lifecycle.cancellation.clone() + } else { + None + }, + provider_name: self + .session + .as_ref() + .and_then(ModelSession::provider_name) + .or(self.provider_name.as_deref()), + model_name: self.session.as_ref().and_then(ModelSession::model_name), + error: None, + } + } + + async fn run_turn_finish( + &mut self, + result: &mut TurnResult, + editable: bool, + failure: Option<&LoopError>, + append_output: bool, + ) -> Result<(), LoopError> { + self.lifecycle.finish_attempted = true; + if self.lifecycle.terminal.is_none() { + self.begin_terminal(result, editable, failure, append_output); + } + loop { + let state = self.lifecycle.terminal.as_ref().expect("terminal state"); + if state.finish_done || state.finish_index == self.mutators.len() { + break; + } + let visible = state.result.clone(); + let mut candidate = visible.clone(); + let output = if state.editable { + Some(output_payload(&visible.items, &mut candidate.items)) + } else { + None + }; + let outcome = self.mutators[state.finish_index] + .on_turn_finish( + TurnFinishPayload { + result: &visible, + output, + }, + LifecycleCtx { + error: state.failure.as_ref(), + ..self.lifecycle_ctx(Some(&visible.turn_id), None) + }, + ) + .await; + let state = self.lifecycle.terminal.as_mut().expect("terminal state"); + state.result = candidate; + state.finish_index += 1; + if let Err(error) = outcome { + state.finish_error.get_or_insert(error); + } + } + let state = self.lifecycle.terminal.as_mut().expect("terminal state"); + if !state.finish_done { + if let Err(error) = validate_output_edits(&state.original_items, &state.result.items) { + state.finish_error.get_or_insert(error); + } + state.finish_done = true; + } + *result = state.result.clone(); + state.finish_error.clone().map_or(Ok(()), Err) + } + + async fn resume_terminal_delivery(&mut self) -> Result { + let state = self.lifecycle.terminal.as_ref().expect("terminal state"); + let mut result = state.result.clone(); + let failure = state.failure.clone(); + let editable = state.editable; + let append_output = !state.output_committed; + let mut first_error = failure.clone(); + if let Err(error) = self + .run_turn_finish(&mut result, editable, failure.as_ref(), append_output) + .await + { + first_error.get_or_insert(error.clone()); + let state = self.lifecycle.terminal.as_mut().expect("terminal state"); + if !state.output_committed { + result.items.clear(); + result.finish_reason = FinishReason::Error; + state.result = result.clone(); + state.output_committed = true; + } + state.failure.get_or_insert(error); + } + if !self + .lifecycle + .terminal + .as_ref() + .expect("terminal state") + .output_committed + { + self.extend_transcript(result.items.clone()); + self.mark_terminal_output_committed(&result.turn_id); + self.close_interrupted_tool_calls(); + } + if let Err(error) = self.run_ready_tool_batches().await { + first_error.get_or_insert(error); + } + let failure = self + .lifecycle + .terminal + .as_ref() + .expect("terminal state") + .failure + .clone(); + if let Err(error) = self.finish_logical_turn(&result, failure.as_ref()).await { + first_error.get_or_insert(error); + } + first_error.map_or(Ok(result), Err) + } + + async fn notify_model_error(&mut self, error: &LoopError, turn_id: &agentkit_core::TurnId) { + for mutator in &self.mutators { + if let Err(hook_error) = mutator + .on_model_error( + error, + self.lifecycle_ctx(Some(turn_id), Some(self.lifecycle.model_call_index)), + ) + .await + { + tracing::debug!(%hook_error, "model error callback failed"); + } + } + } + + async fn run_ready_tool_batches(&mut self) -> Result<(), LoopError> { + let mut index = 0; + let mut first_error = None; + while index < self.pending_tool_batches.len() { + let batch = &self.pending_tool_batches[index]; + if batch.results.len() != batch.calls.len() { + index += 1; + continue; + } + while self.pending_tool_batches[index].next_mutator < self.mutators.len() { + let batch = &self.pending_tool_batches[index]; + let results: Vec = batch + .calls + .iter() + .map(|call| batch.results[&call.id].clone()) + .collect(); + let outcome = self.mutators[batch.next_mutator] + .on_tool_batch( + ToolBatchPayload { + calls: &batch.calls, + results: &results, + }, + self.lifecycle_ctx(Some(&batch.turn_id), None), + ) + .await; + let batch = &mut self.pending_tool_batches[index]; + batch.next_mutator += 1; + if let Err(error) = outcome { + batch.error.get_or_insert(error); + } + } + let batch = self.pending_tool_batches.remove(index); + if let Some(error) = batch.error { + first_error.get_or_insert(error); + } + } + first_error.map_or(Ok(()), Err) + } + + fn record_batch_result(&mut self, result: &ToolResultPart) { + for batch in &mut self.pending_tool_batches { + if batch.calls.iter().any(|call| call.id == result.call_id) { + batch + .results + .entry(result.call_id.clone()) + .or_insert_with(|| result.clone()); + } } } @@ -1985,7 +2603,7 @@ where self.collect_pending_loop_updates().await?; let mut resolutions = std::mem::take(&mut self.pending_loop_updates); if !resolutions.is_empty() { - self.start_logical_turn(); + self.start_logical_turn().await?; } let mut saw_items = false; while let Some(resolution) = resolutions.pop_front() { @@ -1995,7 +2613,7 @@ where saw_items = true; } TaskResolution::Approval(task) => { - let turn_id = self.start_logical_turn(); + let turn_id = self.start_logical_turn().await?; self.enqueue_pending_approval(&turn_id, task, None); } } @@ -2044,8 +2662,9 @@ where session_id: &observed_session_id, observers: &observers, }; + let mut candidate = self.transcript.clone(); let mut cursor = TranscriptCursor { - items: &mut self.transcript, + items: &mut candidate, dirty: false, }; for mutator in &mutators { @@ -2066,6 +2685,7 @@ where } if cursor.dirty { validate_transcript_invariants(cursor.items)?; + self.commit_transcript(candidate); } Ok(()) } @@ -2094,6 +2714,7 @@ where self.active_tool_round = None; return self .finish_cancelled(presentation_turn_id, Vec::new()) + .await .map(Some); } @@ -2144,7 +2765,8 @@ where TaskStartOutcome::Pending { kind, .. } => { self.emit(AgentEvent::ToolExecutionStarted(call.clone())); if kind == agentkit_task_manager::TaskKind::Background { - self.append_detach_placeholder(call.id.clone(), &call.name); + self.append_detach_placeholder(call.id.clone(), &call.name) + .await?; if let Some(active) = self.active_tool_round.as_mut() { active.background_pending = true; } @@ -2179,7 +2801,8 @@ where } } Some(TurnTaskUpdate::Detached(snapshot)) => { - self.append_detach_placeholder(snapshot.call_id, &snapshot.tool_name); + self.append_detach_placeholder(snapshot.call_id, &snapshot.tool_name) + .await?; if let Some(active) = self.active_tool_round.as_mut() { active.background_pending = true; active.foreground_progressed = true; @@ -2199,6 +2822,7 @@ where self.active_tool_round = None; return self .finish_cancelled(presentation_turn_id, Vec::new()) + .await .map(Some); } let active = self.active_tool_round.take().ok_or_else(|| { @@ -2253,17 +2877,33 @@ where turn_id: agentkit_core::TurnId, mutation_point: MutationPoint, ) -> Result { + // Preserve cancellation that arrived while awaiting a logical-start hook. + // Existing approval/tool continuations retain their checkpoint policy. + if mutation_point == MutationPoint::AfterTurnEnded + && self + .lifecycle + .cancellation + .as_ref() + .is_some_and(TurnCancellation::is_cancelled) + { + return self + .finish_cancelled(turn_id, interrupted_assistant_items()) + .await; + } let cancellation = self .cancellation .as_ref() .map(CancellationHandle::checkpoint); + self.lifecycle.cancellation = cancellation.clone(); match self .run_mutators(mutation_point, Some(&turn_id), cancellation.clone()) .await { Ok(()) => {} Err(LoopError::Cancelled) => { - return self.finish_cancelled(turn_id, interrupted_assistant_items()); + return self + .finish_cancelled(turn_id, interrupted_assistant_items()) + .await; } Err(error) => return Err(error), } @@ -2274,14 +2914,16 @@ where // for the model to respond to. Finish the turn rather than dispatch an // assistant-prefill request, which most providers reject. if !transcript_has_pending_input(&self.transcript) { - let turn_result = TurnResult { + let mut turn_result = TurnResult { turn_id, finish_reason: FinishReason::Completed, items: Vec::new(), usage: None, metadata: MetadataMap::new(), }; - self.finish_logical_turn(&turn_result); + self.run_turn_finish(&mut turn_result, true, None, true) + .await?; + self.finish_logical_turn(&turn_result, None).await?; return Ok(LoopStep::Finished(turn_result)); } @@ -2289,13 +2931,15 @@ where .as_ref() .is_some_and(TurnCancellation::is_cancelled) { - return self.finish_cancelled(turn_id, interrupted_assistant_items()); + return self + .finish_cancelled(turn_id, interrupted_assistant_items()) + .await; } let catalog_events = self.tool_executor.drain_catalog_events(); self.emit_tool_catalog_events(catalog_events); - let request = TurnRequest { + let mut request = TurnRequest { session_id: self.session_id.clone(), turn_id: turn_id.clone(), transcript: self.transcript.clone(), @@ -2307,6 +2951,25 @@ where metadata: MetadataMap::new(), }; + self.run_ready_tool_batches().await?; + self.lifecycle.model_call_index += 1; + let original_request = request.transcript.clone(); + for mutator in &self.mutators { + mutator + .on_model_request( + ModelRequestPayload { + transcript: &mut request.transcript, + available_tools: &mut request.available_tools, + cache: &mut request.cache, + metadata: &mut request.metadata, + }, + self.lifecycle_ctx(Some(&turn_id), Some(self.lifecycle.model_call_index)), + ) + .await?; + } + if !self.mutators.is_empty() { + validate_prompt_edits(&original_request, &request.transcript)?; + } let session = self .session .as_mut() @@ -2370,13 +3033,20 @@ where { Ok(turn) => turn, Err(LoopError::Cancelled) => { + self.notify_model_error(&LoopError::Cancelled, &turn_id) + .await; self.task_manager .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled(turn_id, interrupted_assistant_items()); + return self + .finish_cancelled(turn_id, interrupted_assistant_items()) + .await; + } + Err(error) => { + self.notify_model_error(&error, &turn_id).await; + return Err(error); } - Err(error) => return Err(error), }; // begin_turn may apply per-turn routing. Sample the effective selection @@ -2397,7 +3067,6 @@ where } } - let mut saw_tool_call = false; let mut finished_result = None; let mut latest_usage = None; let mut streamed_content = cancellation @@ -2411,20 +3080,39 @@ where { Ok(event) => event, Err(LoopError::Cancelled) => { + self.notify_model_error(&LoopError::Cancelled, &turn_id) + .await; self.task_manager .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled( - turn_id, - interrupted_stream_items(streamed_content.as_ref()), - ); + return self + .finish_cancelled(turn_id, interrupted_stream_items(streamed_content.as_ref())) + .await; + } + Err(error) => { + self.notify_model_error(&error, &turn_id).await; + return Err(error); } - Err(error) => return Err(error), } { + for mutator in &self.mutators { + mutator + .on_model_progress( + &event, + self.lifecycle_ctx(Some(&turn_id), Some(self.lifecycle.model_call_index)), + ) + .await?; + } + for mutator in &self.mutators { + mutator + .on_turn_progress( + TurnProgress::Model(&event), + self.lifecycle_ctx(Some(&turn_id), Some(self.lifecycle.model_call_index)), + ) + .await?; + } let attempt_superseded = matches!(event, ModelTurnEvent::ResponseAttemptSuperseded); if attempt_superseded { - saw_tool_call = false; latest_usage = None; if let Some(content) = &mut streamed_content { content.reset(); @@ -2435,15 +3123,16 @@ where .as_ref() .is_some_and(TurnCancellation::is_cancelled) { + self.notify_model_error(&LoopError::Cancelled, &turn_id) + .await; turn.on_cancelled(); self.task_manager .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled( - turn_id, - interrupted_stream_items(streamed_content.as_ref()), - ); + return self + .finish_cancelled(turn_id, interrupted_stream_items(streamed_content.as_ref())) + .await; } if attempt_superseded { continue; @@ -2460,7 +3149,6 @@ where self.emit(AgentEvent::UsageUpdated(usage)); } ModelTurnEvent::ToolCall(call) => { - saw_tool_call = true; if let Some(content) = &mut streamed_content { content.commit_tool_call(&call); } @@ -2476,10 +3164,31 @@ where } } - let mut result = finished_result.ok_or_else(|| { - LoopError::Provider("model turn ended without a Finished event".into()) - })?; + let Some(mut result) = finished_result else { + let error = LoopError::Provider("model turn ended without a Finished event".into()); + self.notify_model_error(&error, &turn_id).await; + return Err(error); + }; result.usage = merge_usage(result.usage, latest_usage); + let original_output = result.output_items.clone(); + for mutator in &self.mutators { + let visible = result.output_items.clone(); + mutator + .on_model_response( + ModelResponsePayload { + output: output_payload(&visible, &mut result.output_items), + finish_reason: &result.finish_reason, + usage: result.usage.as_ref(), + metadata: &result.metadata, + model: result.model.as_deref(), + response_id: result.response_id.as_deref(), + }, + self.lifecycle_ctx(Some(&turn_id), Some(self.lifecycle.model_call_index)), + ) + .await?; + } + validate_output_edits(&original_output, &result.output_items)?; + let saw_tool_call = !extract_tool_calls(&result.output_items).is_empty(); if let Some(model) = &result.model { chat_span.record("gen_ai.response.model", model.as_str()); } @@ -2506,14 +3215,7 @@ where "gen_ai.response.finish_reasons", provider_finish_reasons(&result.metadata, &result.finish_reason), ); - if let Some(capture) = self.telemetry.output_messages { - record_string_array_attribute( - &chat_span, - "gen_ai.output.messages", - capture_messages(&result.output_items, capture, CaptureOrder::OldestHead), - ); - } - drop(chat_span); + tracing::Span::current().record("saw_tool_call", saw_tool_call); tracing::Span::current().record( "finish_reason", @@ -2540,9 +3242,41 @@ where item }) .collect(); - self.extend_transcript(output_items.clone()); + let mut turn_result = TurnResult { + turn_id: turn_id.clone(), + finish_reason: result.finish_reason.clone(), + items: output_items, + usage: result.usage.clone(), + metadata: result.metadata.clone(), + }; + if !saw_tool_call { + let editable = !matches!( + turn_result.finish_reason, + FinishReason::Cancelled | FinishReason::Error | FinishReason::Blocked + ); + self.run_turn_finish(&mut turn_result, editable, None, true) + .await?; + } + if let Some(capture) = self.telemetry.output_messages { + record_string_array_attribute( + &chat_span, + "gen_ai.output.messages", + capture_messages(&turn_result.items, capture, CaptureOrder::OldestHead), + ); + } + drop(chat_span); + self.extend_transcript(turn_result.items.clone()); + self.mark_terminal_output_committed(&turn_result.turn_id); + let output_items = turn_result.items.clone(); if saw_tool_call { + self.pending_tool_batches.push(PendingToolBatch { + turn_id: turn_id.clone(), + calls: extract_tool_calls(&output_items), + results: BTreeMap::new(), + next_mutator: 0, + error: None, + }); let pending_calls = extract_tool_calls(&output_items) .into_iter() .map(|call| { @@ -2568,13 +3302,10 @@ where if let Some(step) = self.continue_active_tool_round().await? { return Ok(step); } - self.finish_logical_turn(&TurnResult { - turn_id, - finish_reason: result.finish_reason, - items: output_items, - usage: result.usage, - metadata: result.metadata, - }); + self.run_ready_tool_batches().await?; + self.run_turn_finish(&mut turn_result, false, None, false) + .await?; + self.finish_logical_turn(&turn_result, None).await?; return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput( InputRequest { session_id: self.session_id.clone(), @@ -2583,14 +3314,7 @@ where ))); } - let turn_result = TurnResult { - turn_id, - finish_reason: result.finish_reason, - items: output_items, - usage: result.usage, - metadata: result.metadata, - }; - self.finish_logical_turn(&turn_result); + self.finish_logical_turn(&turn_result, None).await?; Ok(LoopStep::Finished(turn_result)) } @@ -2679,7 +3403,8 @@ where self.append_detach_placeholder( pending.call.id.clone(), &pending.call.name, - ); + ) + .await?; } else { self.active_tool_round = Some(ActiveToolRound { presentation_turn_id: pending.presentation_turn_id.clone(), @@ -2725,26 +3450,53 @@ where } } - fn finish_cancelled( + async fn finish_cancelled( + &mut self, + turn_id: agentkit_core::TurnId, + items: Vec, + ) -> Result { + let result = self.finish_cancelled_inner(turn_id, items).await; + if let Err(error) = &result { + self.abort_active_turn(FinishReason::Cancelled, error).await; + } + result + } + + async fn finish_cancelled_inner( &mut self, turn_id: agentkit_core::TurnId, items: Vec, ) -> Result { let pending = self.drain_pending_approval_items(); self.reject_drained_approvals(pending); - self.extend_transcript(items.clone()); - self.close_interrupted_tool_calls(); - let turn_result = TurnResult { + let mut turn_result = TurnResult { turn_id, finish_reason: FinishReason::Cancelled, items, usage: None, metadata: interrupted_metadata("turn"), }; - self.finish_logical_turn(&turn_result); + let now = Timestamp::now(); + for item in &mut turn_result.items { + item.created_at.get_or_insert(now); + if matches!(item.kind, ItemKind::Assistant) { + item.finish_reason.get_or_insert(FinishReason::Cancelled); + } + } + self.run_turn_finish(&mut turn_result, false, None, true) + .await?; + self.extend_transcript(turn_result.items.clone()); + self.mark_terminal_output_committed(&turn_result.turn_id); + self.close_interrupted_tool_calls(); + self.run_ready_tool_batches().await?; + self.finish_logical_turn(&turn_result, None).await?; Ok(LoopStep::Finished(turn_result)) } + /// Buffer new input. With mutators registered, callbacks are awaited by the + /// next `next()` call before dispatch and InputAccepted is emitted only after + /// successful interception. Use `submit_input_async` for awaited acceptance. + /// /// Internal entry point for buffering user input. Reachable only via /// [`InputRequest::submit`] (resolves an `AwaitingInput` interrupt, /// including the very first one after [`Agent::start`]) and @@ -2755,16 +3507,72 @@ where /// after start-up always flows through one of the typed `submit` /// handles. pub fn submit_input(&mut self, input: Vec) -> Result<(), LoopError> { + self.check_input_submission()?; + if self.mutators.is_empty() { + self.accept_input(input); + } else { + self.pending_admission.push_back(input); + } + Ok(()) + } + + /// Await input interception before accepting this batch. On hook failure no + /// rejected input is committed and the active logical turn is unchanged. + /// Batches are processed in submission order. If an earlier queued batch is + /// rejected, this new batch remains queued for subsequent admission. + /// Sync [`Self::submit_input`] remains compatible but queues interception for + /// `next`, before dispatch, rather than awaiting acceptance in the caller. + pub async fn submit_input_async(&mut self, input: Vec) -> Result<(), LoopError> { + self.check_input_submission()?; + self.pending_admission.push_back(input); + self.admit_queued_input().await + } + + fn check_input_submission(&self) -> Result<(), LoopError> { + if self.lifecycle.closed + || self.lifecycle.closing.is_some() + || self.lifecycle.pending_cancel.is_some() + { + return Err(LoopError::InvalidState( + "session is closed or retiring a turn".into(), + )); + } if self.has_pending_interrupts() { return Err(LoopError::InvalidState( "cannot submit input while an interrupt is pending".into(), )); } + Ok(()) + } + + fn accept_input(&mut self, input: Vec) { self.emit(AgentEvent::InputAccepted { session_id: self.session_id.clone(), items: input.clone(), }); self.pending_input.extend(input); + } + + async fn intercept_input(&mut self, input: &mut Vec) -> Result<(), LoopError> { + for mutator in &self.mutators { + mutator + .on_input( + input, + self.lifecycle_ctx(self.lifecycle.active_turn.as_ref(), None), + ) + .await?; + } + Ok(()) + } + + async fn admit_queued_input(&mut self) -> Result<(), LoopError> { + while let Some(mut input) = self.pending_admission.front().cloned() { + let outcome = self.intercept_input(&mut input).await; + // Reject only this batch, preserving unrelated later submissions. + self.pending_admission.pop_front(); + outcome?; + self.accept_input(input); + } Ok(()) } @@ -2773,6 +3581,11 @@ where /// The override is consumed the next time the driver starts a model turn. /// Session-level defaults still apply to later turns. pub fn set_next_turn_cache(&mut self, cache: PromptCacheRequest) -> Result<(), LoopError> { + if self.lifecycle.closed || self.lifecycle.closing.is_some() { + return Err(LoopError::InvalidState( + "session is closed or closing".into(), + )); + } if self.has_pending_interrupts() { return Err(LoopError::InvalidState( "cannot update next-turn cache while an interrupt is pending".into(), @@ -2806,6 +3619,11 @@ where call_id: ToolCallId, decision: ApprovalDecision, ) -> Result<(), LoopError> { + if self.lifecycle.closed || self.lifecycle.closing.is_some() { + return Err(LoopError::InvalidState( + "session is closed or closing".into(), + )); + } let Some(pending) = self.pending_approvals.get_mut(&call_id) else { return Err(LoopError::InvalidState(format!( "no approval request is pending for call {}", @@ -2836,6 +3654,11 @@ where call_id: ToolCallId, input: serde_json::Value, ) -> Result<(), LoopError> { + if self.lifecycle.closed || self.lifecycle.closing.is_some() { + return Err(LoopError::InvalidState( + "session is closed or closing".into(), + )); + } let Some(pending) = self.pending_approvals.get_mut(&call_id) else { return Err(LoopError::InvalidState(format!( "no approval request is pending for call {}", @@ -2875,7 +3698,13 @@ where /// /// This clears the blocking approval and appends an error tool result so /// the transcript remains provider-valid if the host continues the turn. + /// If this retires the turn, terminal hooks run on the next `next()` or `close()`. pub fn cancel_pending_approval_for(&mut self, call_id: ToolCallId) -> Result<(), LoopError> { + if self.lifecycle.closing.is_some() || self.lifecycle.closed { + return Err(LoopError::InvalidState( + "session is closed or closing".into(), + )); + } let Some(pending) = self.drain_pending_approval_for(&call_id) else { return Err(LoopError::InvalidState(format!( "no approval request is pending for call {}", @@ -2885,7 +3714,24 @@ where let turn_id = pending.presentation_turn_id.clone(); self.reject_drained_approvals(vec![pending]); if self.pending_approvals.is_empty() && self.active_tool_round.is_none() { - let _ = self.finish_cancelled(turn_id, Vec::new())?; + if self.mutators.is_empty() { + self.close_interrupted_tool_calls(); + if self.pending_round_resume.as_ref() == Some(&turn_id) { + self.pending_round_resume = None; + } + if self.lifecycle.active_turn.as_ref() == Some(&turn_id) { + self.lifecycle.active_turn = None; + self.emit(AgentEvent::TurnFinished(TurnResult { + turn_id, + finish_reason: FinishReason::Cancelled, + items: Vec::new(), + usage: None, + metadata: interrupted_metadata("turn"), + })); + } + } else { + self.lifecycle.pending_cancel = Some(turn_id); + } } Ok(()) } @@ -2896,6 +3742,11 @@ where /// individual approval prompt. Each pending approval is resolved as denied /// and receives an error tool result so the transcript remains valid. pub async fn cancel_pending_approvals(&mut self) -> Result, LoopError> { + if self.lifecycle.closing.is_some() || self.lifecycle.closed { + return Err(LoopError::InvalidState( + "session is closed or closing".into(), + )); + } if self.pending_approvals.is_empty() { return Ok(None); } @@ -2932,16 +3783,10 @@ where self.reject_drained_approvals(pending); if let Some(error) = cleanup_error { self.close_interrupted_tool_calls(); - self.finish_logical_turn(&TurnResult { - turn_id, - finish_reason: FinishReason::Error, - items: Vec::new(), - usage: None, - metadata: MetadataMap::new(), - }); + self.abort_active_turn(FinishReason::Error, &error).await; return Err(error); } - self.finish_cancelled(turn_id, Vec::new()).map(Some) + self.finish_cancelled(turn_id, Vec::new()).await.map(Some) } /// Retire the active logical turn without resuming model or tool execution. @@ -2961,6 +3806,11 @@ where /// emitting the cancelled terminal event. The turn is retired even on error; /// a failing task manager may leave external tasks running. pub async fn retire_interrupted_turn(&mut self) -> Result, LoopError> { + if self.lifecycle.closing.is_some() || self.lifecycle.closed { + return Err(LoopError::InvalidState( + "session is closed or closing".into(), + )); + } let Some(turn_id) = self.lifecycle.active_turn.clone() else { return Ok(None); }; @@ -2968,15 +3818,25 @@ where // At AfterToolResult the foreground task round has already completed. self.pending_round_resume = None; let cleanup = self.cleanup_interrupted_turn().await; - let result = TurnResult { + let mut result = TurnResult { turn_id, finish_reason: FinishReason::Cancelled, items: Vec::new(), usage: None, metadata: interrupted_metadata("turn"), }; - self.finish_logical_turn(&result); + self.lifecycle.pending_cancel = None; + let finish = self + .run_turn_finish(&mut result, false, cleanup.as_ref().err(), false) + .await; + let batch = self.run_ready_tool_batches().await; + let end = self + .finish_logical_turn(&result, cleanup.as_ref().err()) + .await; cleanup?; + finish?; + batch?; + end?; Ok(Some(result)) } @@ -2985,7 +3845,12 @@ where LoopSnapshot { session_id: self.session_id.clone(), transcript: self.transcript.clone(), - pending_input: self.pending_input.clone(), + pending_input: self + .pending_input + .iter() + .cloned() + .chain(self.pending_admission.iter().flatten().cloned()) + .collect(), } } @@ -3031,45 +3896,240 @@ where /// Returns [`LoopError::InvalidState`] if called while an unresolved /// interrupt is pending, or propagates provider / tool / compaction errors. pub async fn next(&mut self) -> Result { - if self.lifecycle.active_turn.is_none() { - let continuation_turn = self - .pending_approval_order - .iter() - .find_map(|call_id| self.pending_approvals.get(call_id)) - .map(|pending| pending.presentation_turn_id.clone()) - .or_else(|| { - self.active_tool_round - .as_ref() - .map(|active| active.presentation_turn_id.clone()) - }) - .or_else(|| self.pending_round_resume.clone()); - if let Some(turn_id) = continuation_turn { - self.start_logical_turn_with(turn_id); - } else if !self.pending_input.is_empty() { - self.start_logical_turn(); + if self.lifecycle.closed || self.lifecycle.closing.is_some() { + return Err(LoopError::InvalidState( + "session is closed or closing".into(), + )); + } + let result = async { + if self.lifecycle.terminal.is_some() { + return self + .resume_terminal_delivery() + .await + .map(LoopStep::Finished); } + if let Some(turn_id) = self.lifecycle.pending_cancel.take() { + return self.finish_cancelled(turn_id, Vec::new()).await; + } + self.admit_queued_input().await?; + if self.lifecycle.active_turn.is_none() { + let continuation_turn = self + .pending_approval_order + .iter() + .find_map(|call_id| self.pending_approvals.get(call_id)) + .map(|pending| pending.presentation_turn_id.clone()) + .or_else(|| { + self.active_tool_round + .as_ref() + .map(|active| active.presentation_turn_id.clone()) + }) + .or_else(|| self.pending_round_resume.clone()); + if let Some(turn_id) = continuation_turn { + self.start_logical_turn_with(turn_id).await?; + } else if !self.pending_input.is_empty() { + self.start_logical_turn().await?; + } + } + let step = self.next_inner().await?; + self.run_ready_tool_batches().await?; + Ok(step) + } + .await; + if let Err(error) = &result { + let reason = if matches!(error, LoopError::Cancelled) { + FinishReason::Cancelled + } else { + FinishReason::Error + }; + self.abort_active_turn(reason, error).await; } + result + } + + async fn abort_active_turn(&mut self, reason: FinishReason, failure: &LoopError) { + self.lifecycle.pending_cancel = None; + if let Some(turn_id) = self.lifecycle.active_turn.clone() { + if let Err(error) = self.cleanup_interrupted_turn().await { + tracing::debug!(%error, "failed to clean up turn after loop error"); + } + let mut result = TurnResult { + turn_id, + finish_reason: reason, + items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + }; + if let Some(state) = &mut self.lifecycle.terminal { + state.result = result.clone(); + state.output_committed = true; + state.editable = false; + state.failure = Some(failure.clone()); + } + if !self.lifecycle.finish_attempted { + if let Err(error) = self + .run_turn_finish(&mut result, false, Some(failure), false) + .await + { + tracing::debug!(%error, "terminal finish callback failed"); + } + } + if let Err(error) = self.run_ready_tool_batches().await { + tracing::debug!(%error, "terminal batch callback failed"); + } + if let Err(error) = self.finish_logical_turn(&result, Some(failure)).await { + tracing::debug!(%error, "terminal end callback failed"); + } + } + } - let result = self.next_inner().await; - match &result { - Ok(LoopStep::Finished(turn)) => self.finish_logical_turn(turn), - Err(_) => { + /// Explicitly close this session, retiring an active turn as cancelled. + /// If final delivery already began in `next`, its existing candidate, + /// outcome and callback progress are resumed instead of being replaced. + /// Admission and `next` are disabled as soon as closing starts. Dropping this + /// future leaves resumable progress: a later `close` retries the in-flight + /// callback, but never repeats completed callbacks or the terminal event. + /// Callbacks must tolerate cancellation/retry of their own unfinished work. + /// All cleanup/end callbacks are attempted, with the first error returned. + /// After completion (even on error), repeated closes are no-ops. Detached + /// background work is not awaited; Drop does not run callbacks. + pub async fn close(&mut self) -> Result<(), LoopError> { + if self.lifecycle.closed { + return Ok(()); + } + if self.lifecycle.closing.is_none() { + let mut cleanup_turns = Vec::new(); + if let Some(active) = &self.active_tool_round { + cleanup_turns.push(active.task_turn_id.clone()); + } + if let Some(turn_id) = &self.pending_round_resume { + if !cleanup_turns.contains(turn_id) { + cleanup_turns.push(turn_id.clone()); + } + } + for pending in self.pending_approvals.values() { + if !cleanup_turns.contains(&pending.tool_request.turn_id) { + cleanup_turns.push(pending.tool_request.turn_id.clone()); + } + } + if self.lifecycle.terminal.is_none() { if let Some(turn_id) = self.lifecycle.active_turn.clone() { - if let Err(error) = self.cleanup_interrupted_turn().await { - tracing::debug!(%error, "failed to clean up turn after loop error"); - } - self.finish_logical_turn(&TurnResult { + let result = TurnResult { turn_id, - finish_reason: FinishReason::Error, + finish_reason: FinishReason::Cancelled, items: Vec::new(), usage: None, - metadata: MetadataMap::new(), - }); + metadata: interrupted_metadata("turn"), + }; + self.begin_terminal(&result, false, None, false); } } - _ => {} + self.lifecycle.closing = Some(ClosingState { + cleanup_turns, + cleanup_index: 0, + cleanup_done: false, + session_index: 0, + error: None, + }); } - result + loop { + let state = self.lifecycle.closing.as_ref().expect("closing state"); + if state.cleanup_index == state.cleanup_turns.len() { + break; + } + let turn_id = state.cleanup_turns[state.cleanup_index].clone(); + let outcome = self.task_manager.on_turn_interrupted(&turn_id).await; + let state = self.lifecycle.closing.as_mut().expect("closing state"); + state.cleanup_index += 1; + if let Err(error) = outcome { + state + .error + .get_or_insert(LoopError::Tool(ToolError::Internal(error.to_string()))); + } + } + if !self + .lifecycle + .closing + .as_ref() + .expect("closing state") + .cleanup_done + { + self.active_tool_round = None; + self.pending_round_resume = None; + let pending = self.drain_pending_approval_items(); + self.reject_drained_approvals(pending); + self.close_interrupted_tool_calls(); + self.pending_input.clear(); + self.pending_admission.clear(); + self.lifecycle.pending_cancel = None; + self.lifecycle + .closing + .as_mut() + .expect("closing state") + .cleanup_done = true; + } + if self.lifecycle.terminal.is_some() { + let cleanup_error = self + .lifecycle + .closing + .as_ref() + .expect("closing state") + .error + .clone(); + if let Some(error) = cleanup_error { + self.lifecycle + .terminal + .as_mut() + .expect("terminal state") + .failure + .get_or_insert(error); + } + if let Err(error) = self.resume_terminal_delivery().await { + self.lifecycle + .closing + .as_mut() + .expect("closing state") + .error + .get_or_insert(error); + } + } else if let Err(error) = self.run_ready_tool_batches().await { + self.lifecycle + .closing + .as_mut() + .expect("closing state") + .error + .get_or_insert(error); + } + loop { + let state = self.lifecycle.closing.as_ref().expect("closing state"); + if state.session_index == self.mutators.len() { + break; + } + let outcome = self.mutators[state.session_index] + .on_session_end( + SessionEndPayload { + transcript: &self.transcript, + }, + LifecycleCtx { + error: state.error.as_ref(), + ..self.lifecycle_ctx(None, None) + }, + ) + .await; + let state = self.lifecycle.closing.as_mut().expect("closing state"); + state.session_index += 1; + if let Err(error) = outcome { + state.error.get_or_insert(error); + } + } + self.session = None; + self.pending_tool_batches.clear(); + self.lifecycle.closed = true; + self.lifecycle + .closing + .take() + .expect("closing state") + .error + .map_or(Ok(()), Err) } async fn cleanup_interrupted_turn(&mut self) -> Result<(), LoopError> { @@ -3134,11 +4194,14 @@ where // before unrelated background completions so a delayed approval cannot // bind itself to that turn's TurnStarted event. AfterToolResult resumes // remain ordered ahead of fresh input below. - if self.pending_round_resume.is_none() && !self.pending_input.is_empty() { + if self.pending_round_resume.is_none() + && (!self.pending_input.is_empty() || self.lifecycle.started_with_input) + { + self.lifecycle.started_with_input = false; // Take updates now to preserve the driver's once-per-step manager // handoff, but defer presenting them until this input turn ends. self.collect_pending_loop_updates().await?; - let turn_id = self.start_logical_turn(); + let turn_id = self.start_logical_turn().await?; let drained: Vec = std::mem::take(&mut self.pending_input); self.extend_transcript(drained); return self @@ -3172,7 +4235,7 @@ where ))); } - let turn_id = self.start_logical_turn(); + let turn_id = self.start_logical_turn().await?; let drained: Vec = std::mem::take(&mut self.pending_input); self.extend_transcript(drained); self.drive_turn(turn_id, MutationPoint::AfterTurnEnded) @@ -3187,7 +4250,31 @@ where /// registered [`TranscriptObserver`]s. The single mutation point — /// every push to `self.transcript` should funnel through here so /// observers see exactly what landed in the transcript. + fn commit_transcript(&mut self, mut candidate: Vec) { + if candidate.starts_with(&self.transcript) { + let suffix = candidate.split_off(self.transcript.len()); + self.extend_transcript(suffix); + } else { + let now = Timestamp::now(); + for item in &mut candidate { + item.created_at.get_or_insert(now); + } + self.transcript = candidate; + for observer in &self.transcript_observers { + observer.on_transcript_rewrite(TranscriptRewriteEvent { + session_id: &self.session_id, + items: &self.transcript, + }); + } + } + } + fn append_item(&mut self, mut item: Item) { + for part in &item.parts { + if let Part::ToolResult(result) = part { + self.record_batch_result(result); + } + } if item.created_at.is_none() { item.created_at = Some(Timestamp::now()); } @@ -3200,10 +4287,13 @@ where self.transcript.push(item); } - fn append_detach_placeholder(&mut self, call_id: ToolCallId, tool_name: &str) { - self.background_call_ids.insert(call_id.clone()); - if !self.detached_call_ids.insert(call_id.clone()) { - return; + async fn append_detach_placeholder( + &mut self, + call_id: ToolCallId, + tool_name: &str, + ) -> Result<(), LoopError> { + if self.detached_call_ids.contains(&call_id) { + return Ok(()); } let detached_result = ToolResultPart { call_id: call_id.clone(), @@ -3213,6 +4303,24 @@ where is_error: false, metadata: MetadataMap::new(), }; + for mutator in &self.mutators { + mutator + .on_tool_progress( + &detached_result, + self.lifecycle_ctx(self.lifecycle.active_turn.as_ref(), None), + ) + .await?; + } + for mutator in &self.mutators { + mutator + .on_turn_progress( + TurnProgress::ToolDetached(&detached_result), + self.lifecycle_ctx(self.lifecycle.active_turn.as_ref(), None), + ) + .await?; + } + self.background_call_ids.insert(call_id.clone()); + self.detached_call_ids.insert(call_id); self.emit(AgentEvent::ToolExecutionProgress(detached_result.clone())); self.append_item(Item { id: None, @@ -3223,6 +4331,7 @@ where finish_reason: None, created_at: None, }); + Ok(()) } /// Append a tool-result Item: emit one [`AgentEvent::ToolResultReceived`] @@ -4808,7 +5917,7 @@ fn tool_result_not_started(item: &Item) -> bool { } /// Errors that can occur while driving the agent loop. -#[derive(Debug, Error)] +#[derive(Debug, Clone, Error)] pub enum LoopError { /// Typed, sanitized model failure with retry accounting. #[error(transparent)] @@ -4881,6 +5990,102 @@ fn fan_out_observed_event( /// The only invariant currently checked is tool_use ↔ tool_result pairing /// — every [`Part::ToolCall`] must be followed (in transcript order) by a /// matching [`Part::ToolResult`] with the same `call_id`. +fn validate_prompt_edits(original: &[Item], candidate: &[Item]) -> Result<(), LoopError> { + validate_transcript_invariants(candidate)?; + let identified = original.iter().filter(|item| item.id.is_some()); + for original_item in identified { + let matches: Vec<&Item> = candidate + .iter() + .filter(|item| item.id == original_item.id) + .collect(); + if matches.len() != 1 + || matches[0].kind != original_item.kind + || matches[0].usage != original_item.usage + || matches[0].finish_reason != original_item.finish_reason + || matches[0].created_at != original_item.created_at + { + return Err(LoopError::Mutator( + "prompt hook changed a message identity or accounting".into(), + )); + } + } + for original_item in original.iter().filter(|item| { + item.id.is_none() + && (item.usage.is_some() || item.finish_reason.is_some() || item.created_at.is_some()) + }) { + let same_facts = |item: &&Item| { + item.id.is_none() + && item.kind == original_item.kind + && item.usage == original_item.usage + && item.finish_reason == original_item.finish_reason + && item.created_at == original_item.created_at + }; + if candidate.iter().filter(same_facts).count() < original.iter().filter(same_facts).count() + { + return Err(LoopError::Mutator( + "prompt hook removed or changed anonymous item accounting".into(), + )); + } + } + let tool_parts = |items: &[Item]| -> Vec { + items + .iter() + .flat_map(|item| &item.parts) + .filter(|part| matches!(part, Part::ToolCall(_) | Part::ToolResult(_))) + .cloned() + .collect() + }; + if tool_parts(original) != tool_parts(candidate) { + return Err(LoopError::Mutator( + "prompt hook changed existing tool calls or results".into(), + )); + } + Ok(()) +} + +fn output_payload<'a>( + original: &'a [Item], + candidate: &'a mut [Item], +) -> Vec> { + original + .iter() + .zip(candidate) + .map(|(original, item)| OutputItemPayload { + original, + parts: &mut item.parts, + metadata: &mut item.metadata, + }) + .collect() +} + +fn validate_output_edits(original: &[Item], candidate: &[Item]) -> Result<(), LoopError> { + fn links(item: &Item) -> Vec<(bool, ToolCallId, Option)> { + item.parts + .iter() + .filter_map(|part| match part { + Part::ToolCall(call) => Some((true, call.id.clone(), Some(call.name.clone()))), + Part::ToolResult(result) => Some((false, result.call_id.clone(), None)), + _ => None, + }) + .collect() + } + if original.len() != candidate.len() + || original.iter().zip(candidate).any(|(a, b)| { + a.id != b.id + || a.kind != b.kind + || a.usage != b.usage + || a.finish_reason != b.finish_reason + || a.created_at != b.created_at + || links(a) != links(b) + }) + { + return Err(LoopError::Mutator( + "output hook changed immutable item identity, accounting, or tool linkage".into(), + )); + } + Ok(()) +} + fn validate_transcript_invariants(transcript: &[Item]) -> Result<(), LoopError> { let mut pending: HashSet = HashSet::new(); let mut seen_calls: HashSet = HashSet::new(); @@ -8786,6 +9991,7 @@ mod tests { )); driver .finish_cancelled(agentkit_core::TurnId::new("turn-cancelled"), Vec::new()) + .await .unwrap(); assert!( @@ -10091,3 +11297,6 @@ mod tests { ); } } + +#[cfg(test)] +mod lifecycle_tests; diff --git a/crates/agentkit-loop/src/lifecycle_tests.rs b/crates/agentkit-loop/src/lifecycle_tests.rs new file mode 100644 index 0000000..24812bc --- /dev/null +++ b/crates/agentkit-loop/src/lifecycle_tests.rs @@ -0,0 +1,1415 @@ +use super::*; +use agentkit_tools_core::{ToolContext, ToolExecutionOutcome, ToolResult}; +use std::sync::{ + Mutex, + atomic::{AtomicBool, Ordering}, +}; + +#[derive(Clone, Copy)] +enum Script { + Reply, + Tool, + BeginError, + StreamError, + MissingFinished, + Cancelled, +} + +#[derive(Clone)] +struct ProbeModel { + scripts: Arc>>, + requests: Arc>>, + configs: Arc>>, +} +impl ProbeModel { + fn new(scripts: impl IntoIterator) -> Self { + Self { + scripts: Arc::new(Mutex::new(scripts.into_iter().collect())), + requests: Arc::default(), + configs: Arc::default(), + } + } +} +struct ProbeSession(ProbeModel); +struct ProbeTurn(VecDeque>); +#[async_trait] +impl ModelAdapter for ProbeModel { + type Session = ProbeSession; + async fn start_session(&self, config: SessionConfig) -> Result { + self.configs.lock().unwrap().push(config); + Ok(ProbeSession(self.clone())) + } +} +#[async_trait] +impl ModelSession for ProbeSession { + type Turn = ProbeTurn; + async fn begin_turn( + &mut self, + request: TurnRequest, + _: Option, + ) -> Result { + self.0.requests.lock().unwrap().push(request); + let script = self + .0 + .scripts + .lock() + .unwrap() + .pop_front() + .expect("unexpected inference"); + let item = match script { + Script::BeginError => return Err(LoopError::Provider("original begin error".into())), + Script::Cancelled => return Err(LoopError::Cancelled), + Script::StreamError => { + return Ok(ProbeTurn(VecDeque::from([Err(LoopError::Provider( + "original stream error".into(), + ))]))); + } + Script::MissingFinished => return Ok(ProbeTurn(VecDeque::new())), + Script::Reply => Item::text(ItemKind::Assistant, "reply").with_id("message-reply"), + Script::Tool => Item::new( + ItemKind::Assistant, + vec![Part::ToolCall(ToolCallPart::new( + "call-probe", + "probe", + serde_json::json!({"value": 1}), + ))], + ) + .with_id("message-tool"), + }; + let finish_reason = if matches!(script, Script::Tool) { + FinishReason::ToolCall + } else { + FinishReason::Completed + }; + let mut events = VecDeque::new(); + if let Some(Part::ToolCall(call)) = item.parts.first() { + events.push_back(Ok(ModelTurnEvent::ToolCall(call.clone()))); + } + events.push_back(Ok(ModelTurnEvent::Finished(ModelTurnResult { + finish_reason, + output_items: vec![item], + usage: None, + metadata: MetadataMap::new(), + model: Some("probe-model".into()), + response_id: Some("response-probe".into()), + }))); + Ok(ProbeTurn(events)) + } +} +#[async_trait] +impl ModelTurn for ProbeTurn { + async fn next_event( + &mut self, + _: Option, + ) -> Result, LoopError> { + self.0.pop_front().transpose() + } +} + +struct ProbeExecutor; +#[async_trait] +impl ToolExecutor for ProbeExecutor { + fn specs(&self) -> Vec { + Vec::new() + } + async fn execute(&self, request: ToolRequest, _: &mut ToolContext<'_>) -> ToolExecutionOutcome { + ToolExecutionOutcome::Completed(ToolResult { + result: ToolResultPart { + call_id: request.call_id, + output: ToolOutput::Text("tool-result".into()), + is_error: false, + metadata: MetadataMap::new(), + }, + duration: None, + metadata: MetadataMap::new(), + }) + } +} + +#[derive(Clone)] +struct ProbeHook { + label: &'static str, + events: Arc>>, + fail: Option<&'static str>, + failed: Arc, + break_links: bool, +} +impl ProbeHook { + fn new(label: &'static str, events: &Arc>>) -> Self { + Self { + label, + events: events.clone(), + fail: None, + failed: Arc::new(AtomicBool::new(false)), + break_links: false, + } + } + fn event(&self, name: &str, ctx: &LifecycleCtx<'_>) -> Result<(), LoopError> { + self.events.lock().unwrap().push(format!( + "{}:{name}:{}:{}", + self.label, + ctx.turn_id.map(|id| id.0.as_str()).unwrap_or("-"), + ctx.model_call_index + .map(|i| i.to_string()) + .unwrap_or_else(|| "-".into()) + )); + if self.fail == Some(name) && !self.failed.swap(true, Ordering::SeqCst) { + Err(LoopError::Mutator(format!("{} failed {name}", self.label))) + } else { + Ok(()) + } + } + fn rewrite(&self, output: &mut [OutputItemPayload<'_>], suffix: &str) { + for item in output { + for part in item.parts.iter_mut() { + if let Part::Text(text) = part { + text.text.push_str(&format!(":{suffix}{}", self.label)); + } + } + } + } +} +#[async_trait] +impl LoopMutator for ProbeHook { + async fn on_session_start( + &self, + payload: SessionStartPayload<'_>, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + payload.metadata.insert(self.label.into(), true.into()); + self.event("session_start", &ctx) + } + async fn on_session_end( + &self, + _: SessionEndPayload<'_>, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + self.event("session_end", &ctx) + } + async fn on_input( + &self, + input: &mut Vec, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + for item in input { + for part in &mut item.parts { + if let Part::Text(text) = part { + text.text.push_str(self.label); + } + } + } + self.event("input", &ctx) + } + async fn on_turn_start( + &self, + _: &mut TranscriptCursor<'_>, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + self.event("turn_start", &ctx) + } + async fn on_turn_finish( + &self, + mut payload: TurnFinishPayload<'_>, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + if let Some(output) = &mut payload.output { + self.rewrite(output, "f"); + } + if matches!( + payload.result.finish_reason, + FinishReason::Error | FinishReason::Cancelled + ) { + assert!(payload.output.is_none()); + } + self.event("turn_finish", &ctx) + } + async fn on_turn_end(&self, _: &TurnResult, ctx: LifecycleCtx<'_>) -> Result<(), LoopError> { + self.event("turn_end", &ctx) + } + async fn on_model_request( + &self, + payload: ModelRequestPayload<'_>, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + payload.metadata.insert(self.label.into(), true.into()); + self.event("request", &ctx) + } + async fn on_model_response( + &self, + mut payload: ModelResponsePayload<'_>, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + self.rewrite(&mut payload.output, "m"); + if self.break_links { + for item in payload.output { + for part in item.parts.iter_mut() { + if let Part::ToolCall(call) = part { + call.id = ToolCallId::new("forged"); + } + } + } + } + self.event("response", &ctx) + } + async fn on_model_progress( + &self, + _: &ModelTurnEvent, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + self.event("progress", &ctx) + } + async fn on_model_error(&self, _: &LoopError, ctx: LifecycleCtx<'_>) -> Result<(), LoopError> { + self.event("model_error", &ctx) + } + async fn on_tool_batch( + &self, + payload: ToolBatchPayload<'_>, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + assert_eq!(payload.calls.len(), 1); + assert_eq!(payload.results.len(), 1); + assert_eq!(payload.calls[0].id, payload.results[0].call_id); + self.event("batch", &ctx) + } +} +struct Capture(Arc>>); +impl LoopObserver for Capture { + fn handle_event(&self, event: ObservedEvent) { + self.0.lock().unwrap().push(event.event); + } +} +fn count(events: &Arc>>, name: &str) -> usize { + events + .lock() + .unwrap() + .iter() + .filter(|event| event.contains(&format!(":{name}:"))) + .count() +} +fn text(item: &Item) -> &str { + match &item.parts[0] { + Part::Text(text) => &text.text, + _ => panic!("expected text"), + } +} +fn finished(step: LoopStep) -> TurnResult { + match step { + LoopStep::Finished(result) => result, + other => panic!("expected finished, got {other:?}"), + } +} + +#[tokio::test] +async fn lifecycle_sequential_payloads_and_committed_final_output() { + let events = Arc::default(); + let observed: Arc>> = Arc::default(); + let model = ProbeModel::new([Script::Reply]); + let agent = Agent::builder() + .model(model.clone()) + .input(vec![Item::text(ItemKind::User, "input")]) + .mutator(ProbeHook::new("1", &events)) + .mutator(ProbeHook::new("2", &events)) + .observer(Capture(observed.clone())) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("lifecycle")).await.unwrap(); + assert_eq!(model.configs.lock().unwrap()[0].metadata.len(), 2); + let result = finished(driver.next().await.unwrap()); + let request = &model.requests.lock().unwrap()[0]; + assert_eq!(request.session_id, SessionId::new("lifecycle")); + assert_eq!(text(&request.transcript[0]), "input12"); + assert_eq!(request.metadata.len(), 2); + assert_eq!(text(&result.items[0]), "reply:m1:m2:f1:f2"); + assert_eq!(driver.snapshot().transcript.last(), result.items.last()); + assert!( + observed + .lock() + .unwrap() + .iter() + .any(|event| matches!(event, AgentEvent::TurnFinished(turn) if turn == &result)) + ); + let calls = events.lock().unwrap(); + for pair in calls.chunks_exact(2) { + assert!(pair[0].starts_with("1:")); + assert_eq!(&pair[0][1..], &pair[1][1..]); + } + drop(calls); + assert_eq!(count(&events, "turn_start"), 2); + assert_eq!(count(&events, "turn_finish"), 2); + assert_eq!(count(&events, "turn_end"), 2); + assert_eq!(count(&events, "session_end"), 0); + driver.close().await.unwrap(); + driver.close().await.unwrap(); + assert_eq!(count(&events, "session_end"), 2); + assert!(matches!( + driver.next().await, + Err(LoopError::InvalidState(_)) + )); + assert!(matches!( + driver.submit_input(vec![]), + Err(LoopError::InvalidState(_)) + )); + assert!(matches!( + driver.submit_input_async(vec![]).await, + Err(LoopError::InvalidState(_)) + )); +} + +#[tokio::test] +async fn sync_admission_is_pre_dispatch_and_async_admission_is_pre_acceptance() { + let events = Arc::default(); + let observed: Arc>> = Arc::default(); + let agent = Agent::builder() + .model(ProbeModel::new([Script::Reply, Script::Reply])) + .mutator(ProbeHook::new("1", &events)) + .observer(Capture(observed.clone())) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("input")).await.unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "sync")]) + .unwrap(); + assert_eq!(count(&events, "input"), 0); + assert_eq!(driver.snapshot().pending_input.len(), 1); + assert!( + !observed + .lock() + .unwrap() + .iter() + .any(|event| matches!(event, AgentEvent::InputAccepted { .. })) + ); + finished(driver.next().await.unwrap()); + assert_eq!(count(&events, "input"), 1); + driver + .submit_input_async(vec![Item::text(ItemKind::User, "async")]) + .await + .unwrap(); + assert_eq!(count(&events, "input"), 2); + finished(driver.next().await.unwrap()); + assert_eq!(count(&events, "input"), 2); +} + +#[tokio::test] +async fn input_failure_drops_admission_without_starting_a_turn_and_can_recover() { + for awaited in [false, true] { + let events = Arc::default(); + let model = ProbeModel::new([Script::Reply]); + let mut hook = ProbeHook::new("1", &events); + hook.fail = Some("input"); + let agent = Agent::builder() + .model(model.clone()) + .mutator(hook) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("input-failure")) + .await + .unwrap(); + let input = vec![Item::text(ItemKind::User, "rejected")]; + if awaited { + assert!(driver.submit_input_async(input).await.is_err()); + } else { + driver.submit_input(input).unwrap(); + assert!(driver.next().await.is_err()); + } + assert!(driver.snapshot().transcript.is_empty()); + assert!(driver.snapshot().pending_input.is_empty()); + assert_eq!(count(&events, "turn_start"), 0); + assert!(model.requests.lock().unwrap().is_empty()); + driver + .submit_input_async(vec![Item::text(ItemKind::User, "accepted")]) + .await + .unwrap(); + finished(driver.next().await.unwrap()); + } +} + +#[tokio::test] +async fn tool_continuation_keeps_one_logical_turn_and_reports_real_batch() { + let events = Arc::default(); + let model = ProbeModel::new([Script::Tool, Script::Reply]); + let agent = Agent::builder() + .model(model.clone()) + .tool_executor(ProbeExecutor) + .mutator(ProbeHook::new("1", &events)) + .input(vec![Item::text(ItemKind::User, "tools")]) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("tools")).await.unwrap(); + let turn_id = match driver.next().await.unwrap() { + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info.turn_id, + other => panic!("{other:?}"), + }; + assert_eq!(count(&events, "batch"), 1); + assert_eq!(count(&events, "turn_end"), 0); + driver + .submit_input(vec![Item::text(ItemKind::User, "interject")]) + .unwrap(); + let result = finished(driver.next().await.unwrap()); + assert_eq!(result.turn_id, turn_id); + assert_eq!(count(&events, "turn_start"), 1); + assert_eq!(count(&events, "turn_finish"), 1); + assert_eq!(count(&events, "turn_end"), 1); + let requests = model.requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].turn_id, requests[1].turn_id); + let calls = events.lock().unwrap(); + assert!( + calls + .iter() + .any(|event| event.ends_with(":1") && event.contains(":request:")) + ); + assert!( + calls + .iter() + .any(|event| event.ends_with(":2") && event.contains(":request:")) + ); +} + +#[tokio::test] +async fn output_hook_failure_is_atomic_and_finishes_once() { + for point in [ + "response", + "turn_finish", + "progress", + "request", + "turn_start", + ] { + let events = Arc::default(); + let observed: Arc>> = Arc::default(); + let mut hook = ProbeHook::new("1", &events); + hook.fail = Some(point); + let agent = Agent::builder() + .model(ProbeModel::new([Script::Reply])) + .mutator(hook) + .input(vec![Item::text(ItemKind::User, "input")]) + .observer(Capture(observed.clone())) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("atomic")).await.unwrap(); + assert!(matches!(driver.next().await, Err(LoopError::Mutator(_)))); + assert!( + !driver + .snapshot() + .transcript + .iter() + .any(|item| item.kind == ItemKind::Assistant) + ); + assert_eq!(count(&events, "turn_finish"), 1, "{point}"); + assert_eq!(count(&events, "turn_end"), 1, "{point}"); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) + )); + assert_eq!(count(&events, "turn_end"), 1); + assert_eq!( + observed + .lock() + .unwrap() + .iter() + .filter(|event| matches!(event, AgentEvent::TurnFinished(_))) + .count(), + 1 + ); + } +} + +#[tokio::test] +async fn model_failures_preserve_original_error_despite_terminal_hook_failures() { + for script in [ + Script::BeginError, + Script::StreamError, + Script::MissingFinished, + ] { + let events = Arc::default(); + let mut hook = ProbeHook::new("1", &events); + hook.fail = Some("model_error"); + let mut terminal = ProbeHook::new("2", &events); + terminal.fail = Some("turn_finish"); + let mut end = ProbeHook::new("3", &events); + end.fail = Some("turn_end"); + let agent = Agent::builder() + .model(ProbeModel::new([script])) + .mutator(hook) + .mutator(terminal) + .mutator(end) + .input(vec![Item::text(ItemKind::User, "input")]) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("failure")).await.unwrap(); + assert!(matches!(driver.next().await, Err(LoopError::Provider(_)))); + assert_eq!(count(&events, "model_error"), 3); + assert_eq!(count(&events, "turn_finish"), 3); + assert_eq!(count(&events, "turn_end"), 3); + driver.close().await.unwrap(); + assert_eq!(count(&events, "turn_end"), 3); + } +} + +#[tokio::test] +async fn cancellation_and_explicit_retirement_run_readonly_finish_and_end() { + let events = Arc::default(); + let agent = Agent::builder() + .model(ProbeModel::new([Script::Cancelled])) + .mutator(ProbeHook::new("1", &events)) + .input(vec![Item::text(ItemKind::User, "input")]) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("cancelled")).await.unwrap(); + let result = finished(driver.next().await.unwrap()); + assert_eq!(result.finish_reason, FinishReason::Cancelled); + assert_eq!(driver.snapshot().transcript.last(), result.items.last()); + assert_eq!(count(&events, "turn_finish"), 1); + assert_eq!(count(&events, "turn_end"), 1); + assert_eq!(count(&events, "model_error"), 1); + assert!(driver.retire_interrupted_turn().await.unwrap().is_none()); + driver.close().await.unwrap(); + assert_eq!(count(&events, "turn_end"), 1); + + let events = Arc::default(); + let agent = Agent::builder() + .model(ProbeModel::new([Script::Tool])) + .tool_executor(ProbeExecutor) + .mutator(ProbeHook::new("1", &events)) + .input(vec![Item::text(ItemKind::User, "input")]) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("retire")).await.unwrap(); + driver.next().await.unwrap(); + driver.close().await.unwrap(); + assert_eq!(count(&events, "turn_finish"), 1); + assert_eq!(count(&events, "turn_end"), 1); + assert_eq!(count(&events, "session_end"), 1); +} + +#[tokio::test] +async fn forged_tool_linkage_is_rejected_before_commit_or_execution() { + let events = Arc::default(); + let mut hook = ProbeHook::new("1", &events); + hook.break_links = true; + let agent = Agent::builder() + .model(ProbeModel::new([Script::Tool])) + .tool_executor(ProbeExecutor) + .mutator(hook) + .input(vec![Item::text(ItemKind::User, "input")]) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("identity")).await.unwrap(); + assert!(matches!(driver.next().await, Err(LoopError::Mutator(_)))); + assert_eq!(driver.snapshot().transcript.len(), 1); + assert_eq!(count(&events, "batch"), 0); + assert_eq!(count(&events, "turn_end"), 1); +} + +#[tokio::test] +async fn session_end_failure_does_not_skip_later_hooks_or_reopen_session() { + let events = Arc::default(); + let mut hook = ProbeHook::new("1", &events); + hook.fail = Some("session_end"); + let agent = Agent::builder() + .model(ProbeModel::new([])) + .mutator(hook) + .mutator(ProbeHook::new("2", &events)) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("close-error")) + .await + .unwrap(); + assert!(driver.close().await.is_err()); + assert_eq!(count(&events, "session_end"), 2); + driver.close().await.unwrap(); + assert_eq!(count(&events, "session_end"), 2); + assert!(driver.next().await.is_err()); +} + +struct PromptHook { + fail_start: bool, + fail_request: bool, +} +#[async_trait] +impl LoopMutator for PromptHook { + async fn on_turn_start( + &self, + cursor: &mut TranscriptCursor<'_>, + _: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + for item in cursor.iter_mut() { + if item.kind == ItemKind::User { + item.parts = vec![Part::text("logical prompt")]; + } + } + cursor.insert(0, Item::text(ItemKind::System, "persistent context")); + if self.fail_start { + Err(LoopError::Mutator("start prompt failed".into())) + } else { + Ok(()) + } + } + async fn on_model_request( + &self, + payload: ModelRequestPayload<'_>, + _: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + for item in payload.transcript.iter_mut() { + if item.kind == ItemKind::User { + item.parts = vec![Part::text("inference prompt")]; + } + } + payload + .transcript + .insert(0, Item::text(ItemKind::System, "inference-only context")); + if self.fail_request { + Err(LoopError::Mutator("request prompt failed".into())) + } else { + Ok(()) + } + } +} + +#[tokio::test] +async fn logical_prompt_edits_persist_but_inference_prompt_edits_do_not() { + let model = ProbeModel::new([Script::Reply]); + let agent = Agent::builder() + .model(model.clone()) + .mutator(PromptHook { + fail_start: false, + fail_request: false, + }) + .input(vec![ + Item::text(ItemKind::User, "initial").with_id("input-id"), + ]) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("prompts")).await.unwrap(); + finished(driver.next().await.unwrap()); + let requests = model.requests.lock().unwrap(); + assert_eq!(text(&requests[0].transcript[2]), "inference prompt"); + assert_eq!( + requests[0].transcript[2].id, + Some(agentkit_core::MessageId::new("input-id")) + ); + assert_eq!(text(&requests[0].transcript[0]), "inference-only context"); + let transcript = driver.snapshot().transcript; + assert_eq!(text(&transcript[1]), "logical prompt"); + assert_eq!(text(&transcript[0]), "persistent context"); + assert!(!transcript.iter().any(|item| { + item.parts + .iter() + .any(|part| matches!(part, Part::Text(text) if text.text.contains("inference"))) + })); +} + +#[tokio::test] +async fn failed_logical_prompt_edits_never_commit_and_request_edits_never_persist() { + for fail_start in [false, true] { + let model = ProbeModel::new([Script::Reply]); + let agent = Agent::builder() + .model(model.clone()) + .mutator(PromptHook { + fail_start, + fail_request: !fail_start, + }) + .input(vec![Item::text(ItemKind::User, "initial")]) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("prompt-failure")) + .await + .unwrap(); + assert!(matches!(driver.next().await, Err(LoopError::Mutator(_)))); + let transcript = driver.snapshot().transcript; + if fail_start { + assert!(transcript.is_empty()); + } else { + assert_eq!(transcript.len(), 2); + assert_eq!(text(&transcript[1]), "logical prompt"); + } + assert!(model.requests.lock().unwrap().is_empty()); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) + )); + } +} + +#[derive(Default)] +struct StoredTranscript { + items: Mutex>, + appends: Mutex>, + rewrites: std::sync::atomic::AtomicUsize, +} +impl TranscriptObserver for Arc { + fn on_transcript_event(&self, event: TranscriptEvent<'_>) { + self.items.lock().unwrap().push(event.item.clone()); + self.appends.lock().unwrap().push(event.item.clone()); + } + fn on_transcript_rewrite(&self, event: TranscriptRewriteEvent<'_>) { + *self.items.lock().unwrap() = event.items.to_vec(); + self.rewrites.fetch_add(1, Ordering::SeqCst); + } +} + +#[tokio::test] +async fn second_turn_prompt_insertion_emits_rewrite_not_historic_appends() { + let stored = Arc::new(StoredTranscript::default()); + let agent = Agent::builder() + .model(ProbeModel::new([Script::Reply, Script::Reply])) + .mutator(PromptHook { + fail_start: false, + fail_request: false, + }) + .transcript_observer(stored.clone()) + .input(vec![Item::text(ItemKind::User, "first")]) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("rewrite")).await.unwrap(); + finished(driver.next().await.unwrap()); + assert_eq!(*stored.items.lock().unwrap(), driver.snapshot().transcript); + let before = stored.appends.lock().unwrap().len(); + driver + .submit_input_async(vec![Item::text(ItemKind::User, "second")]) + .await + .unwrap(); + finished(driver.next().await.unwrap()); + assert_eq!(*stored.items.lock().unwrap(), driver.snapshot().transcript); + assert_eq!(stored.rewrites.load(Ordering::SeqCst), 1); + assert_eq!( + stored.appends.lock().unwrap().len(), + before + 1, + "only the final output was appended; history/input arrived in the rewrite" + ); +} + +struct RejectInput; +#[async_trait] +impl LoopMutator for RejectInput { + async fn on_input(&self, input: &mut Vec, _: LifecycleCtx<'_>) -> Result<(), LoopError> { + if input.iter().any(|item| text(item) == "reject") { + Err(LoopError::Mutator("rejected batch".into())) + } else { + Ok(()) + } + } +} + +#[tokio::test] +async fn mixed_submission_rejection_preserves_earlier_valid_batch_and_order() { + let model = ProbeModel::new([Script::Reply]); + let agent = Agent::builder() + .model(model.clone()) + .mutator(RejectInput) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("mixed-input")) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "A")]) + .unwrap(); + assert!( + driver + .submit_input_async(vec![Item::text(ItemKind::User, "reject")]) + .await + .is_err() + ); + assert_eq!( + driver + .snapshot() + .pending_input + .iter() + .map(text) + .collect::>(), + ["A"] + ); + driver + .submit_input(vec![Item::text(ItemKind::User, "B")]) + .unwrap(); + driver + .submit_input_async(vec![Item::text(ItemKind::User, "C")]) + .await + .unwrap(); + assert_eq!( + driver + .snapshot() + .pending_input + .iter() + .map(text) + .collect::>(), + ["A", "B", "C"] + ); + finished(driver.next().await.unwrap()); + assert_eq!( + model.requests.lock().unwrap()[0] + .transcript + .iter() + .map(text) + .collect::>(), + ["A", "B", "C"] + ); +} + +#[tokio::test] +async fn queued_rejection_preserves_prior_accepted_and_later_unrelated_batches() { + let model = ProbeModel::new([Script::Reply]); + let agent = Agent::builder() + .model(model.clone()) + .mutator(RejectInput) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("queued-input")) + .await + .unwrap(); + for value in ["A", "reject", "C"] { + driver + .submit_input(vec![Item::text(ItemKind::User, value)]) + .unwrap(); + } + assert!(driver.next().await.is_err()); + assert_eq!( + driver + .snapshot() + .pending_input + .iter() + .map(text) + .collect::>(), + ["A", "C"] + ); + finished(driver.next().await.unwrap()); + assert_eq!( + model.requests.lock().unwrap()[0] + .transcript + .iter() + .map(text) + .collect::>(), + ["A", "C"] + ); +} + +#[derive(Clone)] +struct Gate { + open: Arc, + attempts: Arc, + completed: Arc, +} +impl Gate { + fn new() -> Self { + Self { + open: Arc::new(AtomicBool::new(false)), + attempts: Arc::default(), + completed: Arc::default(), + } + } + async fn wait(&self) { + self.attempts.fetch_add(1, Ordering::SeqCst); + if !self.open.load(Ordering::SeqCst) { + std::future::pending::<()>().await; + } + self.completed.fetch_add(1, Ordering::SeqCst); + } +} +struct BlockingHook { + phase: &'static str, + gate: Gate, +} +#[async_trait] +impl LoopMutator for BlockingHook { + async fn on_turn_finish( + &self, + _: TurnFinishPayload<'_>, + _: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + if self.phase == "finish" { + self.gate.wait().await; + } + Ok(()) + } + async fn on_turn_end(&self, _: &TurnResult, _: LifecycleCtx<'_>) -> Result<(), LoopError> { + if self.phase == "end" { + self.gate.wait().await; + } + Ok(()) + } + async fn on_session_end( + &self, + _: SessionEndPayload<'_>, + _: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + if self.phase == "session" { + self.gate.wait().await; + } + Ok(()) + } + async fn on_tool_batch( + &self, + _: ToolBatchPayload<'_>, + _: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + if self.phase == "batch" { + self.gate.wait().await; + } + Ok(()) + } +} +struct CompactCompletedTools; +#[async_trait] +impl LoopMutator for CompactCompletedTools { + async fn mutate( + &self, + cursor: &mut TranscriptCursor<'_>, + ctx: LoopCtx<'_>, + ) -> Result<(), LoopError> { + if ctx.point == MutationPoint::AfterToolResult { + cursor.retain(|item| { + !item + .parts + .iter() + .any(|part| matches!(part, Part::ToolCall(_) | Part::ToolResult(_))) + }); + } + Ok(()) + } +} + +#[tokio::test] +async fn batch_facts_and_completed_callbacks_survive_compaction_and_dropped_delivery() { + let events = Arc::default(); + let gate = Gate::new(); + let stored = Arc::new(StoredTranscript::default()); + let agent = Agent::builder() + .model(ProbeModel::new([Script::Tool, Script::Reply])) + .tool_executor(ProbeExecutor) + .mutator(ProbeHook::new("first", &events)) + .mutator(BlockingHook { + phase: "batch", + gate: gate.clone(), + }) + .mutator(ProbeHook::new("last", &events)) + .mutator(CompactCompletedTools) + .transcript_observer(stored.clone()) + .input(vec![Item::text(ItemKind::User, "input")]) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("captured-batch")) + .await + .unwrap(); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(10), driver.next()) + .await + .is_err() + ); + assert_eq!(gate.attempts.load(Ordering::SeqCst), 1); + gate.open.store(true, Ordering::SeqCst); + finished(driver.next().await.unwrap()); + assert_eq!(gate.completed.load(Ordering::SeqCst), 1); + assert_eq!( + count(&events, "batch"), + 2, + "each completed first/last callback runs once" + ); + assert!( + !driver + .snapshot() + .transcript + .iter() + .flat_map(|item| &item.parts) + .any(|part| matches!(part, Part::ToolCall(_) | Part::ToolResult(_))) + ); + assert_eq!(*stored.items.lock().unwrap(), driver.snapshot().transcript); + driver.close().await.unwrap(); + assert_eq!(count(&events, "batch"), 2); +} + +struct CleanupManager { + inner: SimpleTaskManager, + gate: Gate, +} +#[async_trait] +impl TaskManager for CleanupManager { + async fn start_task( + &self, + request: TaskLaunchRequest, + ctx: TaskStartContext, + ) -> Result { + self.inner.start_task(request, ctx).await + } + async fn wait_for_turn( + &self, + id: &agentkit_core::TurnId, + cancel: Option, + ) -> Result, agentkit_task_manager::TaskManagerError> { + self.inner.wait_for_turn(id, cancel).await + } + async fn take_pending_loop_updates( + &self, + ) -> Result { + self.inner.take_pending_loop_updates().await + } + async fn on_turn_interrupted( + &self, + id: &agentkit_core::TurnId, + ) -> Result<(), agentkit_task_manager::TaskManagerError> { + self.gate.wait().await; + self.inner.on_turn_interrupted(id).await + } + fn handle(&self) -> agentkit_task_manager::TaskManagerHandle { + self.inner.handle() + } +} + +#[tokio::test] +async fn dropped_close_resumes_cleanup_and_terminal_callbacks_without_duplicates() { + for phase in ["cleanup", "finish", "end", "session"] { + let events = Arc::default(); + let observed: Arc>> = Arc::default(); + let gate = Gate::new(); + let cleanup_gate = Gate::new(); + if phase != "cleanup" { + cleanup_gate.open.store(true, Ordering::SeqCst); + } + let agent = Agent::builder() + .model(ProbeModel::new([Script::Tool])) + .tool_executor(ProbeExecutor) + .task_manager(CleanupManager { + inner: SimpleTaskManager::new(), + gate: cleanup_gate.clone(), + }) + .mutator(ProbeHook::new("first", &events)) + .mutator(BlockingHook { + phase, + gate: gate.clone(), + }) + .mutator(ProbeHook::new("last", &events)) + .observer(Capture(observed.clone())) + .input(vec![Item::text(ItemKind::User, "input")]) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("resumable-close")) + .await + .unwrap(); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) + )); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(10), driver.close()) + .await + .is_err(), + "{phase}" + ); + assert!(matches!( + driver.next().await, + Err(LoopError::InvalidState(_)) + )); + assert!( + driver + .submit_input(vec![Item::text(ItemKind::User, "closed")]) + .is_err() + ); + gate.open.store(true, Ordering::SeqCst); + cleanup_gate.open.store(true, Ordering::SeqCst); + driver.close().await.unwrap(); + driver.close().await.unwrap(); + assert_eq!(cleanup_gate.completed.load(Ordering::SeqCst), 1, "{phase}"); + if phase != "cleanup" { + assert_eq!(gate.completed.load(Ordering::SeqCst), 1, "{phase}"); + } + assert_eq!(count(&events, "turn_finish"), 2, "{phase}"); + assert_eq!(count(&events, "turn_end"), 2, "{phase}"); + assert_eq!(count(&events, "session_end"), 2, "{phase}"); + assert_eq!( + observed + .lock() + .unwrap() + .iter() + .filter(|event| matches!(event, AgentEvent::TurnFinished(_))) + .count(), + 1, + "{phase}" + ); + } +} + +struct InterruptProgress { + controller: agentkit_core::CancellationController, + terminals: Arc, +} +#[async_trait] +impl LoopMutator for InterruptProgress { + async fn on_model_progress( + &self, + _: &ModelTurnEvent, + _: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + self.controller.interrupt(); + Ok(()) + } + async fn on_turn_finish( + &self, + payload: TurnFinishPayload<'_>, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + assert_eq!(payload.result.finish_reason, FinishReason::Cancelled); + assert!(payload.output.is_none()); + assert!(ctx.cancellation.as_ref().unwrap().is_cancelled()); + self.terminals.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + async fn on_turn_end( + &self, + result: &TurnResult, + ctx: LifecycleCtx<'_>, + ) -> Result<(), LoopError> { + assert_eq!(result.finish_reason, FinishReason::Cancelled); + assert!(ctx.cancellation.as_ref().unwrap().is_cancelled()); + self.terminals.fetch_add(1, Ordering::SeqCst); + Ok(()) + } +} +#[tokio::test] +async fn terminal_context_keeps_the_cancelled_active_checkpoint() { + let controller = agentkit_core::CancellationController::new(); + let terminals = Arc::default(); + let agent = Agent::builder() + .model(ProbeModel::new([Script::Reply])) + .cancellation(controller.handle()) + .mutator(InterruptProgress { + controller, + terminals: Arc::clone(&terminals), + }) + .input(vec![Item::text(ItemKind::User, "cancel during progress")]) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("checkpoint")).await.unwrap(); + let result = finished(driver.next().await.unwrap()); + assert_eq!(result.finish_reason, FinishReason::Cancelled); + assert_eq!(terminals.load(Ordering::SeqCst), 2); + assert_eq!(driver.snapshot().transcript.last(), result.items.last()); +} + +#[tokio::test] +async fn readonly_end_notification_failure_preserves_committed_result() { + let events = Arc::default(); + let observed: Arc>> = Arc::default(); + let mut hook = ProbeHook::new("1", &events); + hook.fail = Some("turn_end"); + let agent = Agent::builder() + .model(ProbeModel::new([Script::Reply])) + .mutator(hook) + .observer(Capture(observed.clone())) + .input(vec![Item::text(ItemKind::User, "input")]) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("readonly-end")) + .await + .unwrap(); + let result = finished(driver.next().await.unwrap()); + assert_eq!(result.finish_reason, FinishReason::Completed); + assert_eq!(driver.snapshot().transcript.last(), result.items.last()); + assert!( + observed + .lock() + .unwrap() + .iter() + .any(|event| matches!(event, AgentEvent::TurnFinished(turn) if turn == &result)) + ); + assert_eq!(count(&events, "turn_end"), 1); + driver.close().await.unwrap(); + assert_eq!(count(&events, "turn_end"), 1); +} + +struct SendOnlyModel(ProbeModel); +struct SendOnlySession { + inner: ProbeSession, + calls: std::cell::Cell, +} +struct SendOnlyTurn { + inner: ProbeTurn, + polls: std::cell::Cell, +} +#[async_trait] +impl ModelAdapter for SendOnlyModel { + type Session = SendOnlySession; + async fn start_session(&self, config: SessionConfig) -> Result { + Ok(SendOnlySession { + inner: self.0.start_session(config).await?, + calls: std::cell::Cell::new(0), + }) + } +} +#[async_trait] +impl ModelSession for SendOnlySession { + type Turn = SendOnlyTurn; + async fn begin_turn( + &mut self, + request: TurnRequest, + cancel: Option, + ) -> Result { + self.calls.set(self.calls.get() + 1); + Ok(SendOnlyTurn { + inner: self.inner.begin_turn(request, cancel).await?, + polls: std::cell::Cell::new(0), + }) + } + fn provider_name(&self) -> Option<&str> { + Some("send-only-provider") + } + fn model_name(&self) -> Option<&str> { + Some("send-only-model") + } +} +#[async_trait] +impl ModelTurn for SendOnlyTurn { + async fn next_event( + &mut self, + cancel: Option, + ) -> Result, LoopError> { + self.polls.set(self.polls.get() + 1); + self.inner.next_event(cancel).await + } +} +fn require_send(value: T) -> T { + value +} + +#[tokio::test] +async fn public_lifecycle_futures_are_send_with_non_sync_session_and_turn() { + let events = Arc::default(); + let agent = Agent::builder() + .model(SendOnlyModel(ProbeModel::new([ + Script::Tool, + Script::Reply, + Script::BeginError, + ]))) + .tool_executor(ProbeExecutor) + .mutator(ProbeHook::new("send", &events)) + .build() + .unwrap(); + let mut driver = require_send(agent.start(SessionConfig::new("send-only"))) + .await + .unwrap(); + require_send(driver.submit_input_async(vec![Item::text(ItemKind::User, "tools")])) + .await + .unwrap(); + assert!(matches!( + require_send(driver.next()).await.unwrap(), + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) + )); + finished(require_send(driver.next()).await.unwrap()); + require_send(driver.submit_input_async(vec![Item::text(ItemKind::User, "error")])) + .await + .unwrap(); + assert!(matches!( + require_send(driver.next()).await, + Err(LoopError::Provider(_)) + )); + require_send(driver.close()).await.unwrap(); +} + +#[tokio::test] +async fn earlier_queued_rejection_retains_new_async_input() { + let model = ProbeModel::new([Script::Reply]); + let agent = Agent::builder() + .model(model.clone()) + .mutator(RejectInput) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("inverse-mixed-input")) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "reject")]) + .unwrap(); + assert!( + driver + .submit_input_async(vec![Item::text(ItemKind::User, "B")]) + .await + .is_err() + ); + assert_eq!( + driver + .snapshot() + .pending_input + .iter() + .map(text) + .collect::>(), + ["B"] + ); + finished(driver.next().await.unwrap()); + assert_eq!( + model.requests.lock().unwrap()[0] + .transcript + .iter() + .map(text) + .collect::>(), + ["B"] + ); +} + +#[tokio::test] +async fn dropped_next_terminal_delivery_is_resumed_by_close() { + for phase in ["finish", "end"] { + let events = Arc::default(); + let gate = Gate::new(); + let observed: Arc>> = Arc::default(); + let agent = Agent::builder() + .model(ProbeModel::new([Script::Reply])) + .mutator(ProbeHook::new("first", &events)) + .mutator(BlockingHook { + phase, + gate: gate.clone(), + }) + .mutator(ProbeHook::new("last", &events)) + .observer(Capture(observed.clone())) + .input(vec![Item::text(ItemKind::User, "input")]) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("resume-next-terminal")) + .await + .unwrap(); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(10), driver.next()) + .await + .is_err(), + "{phase}" + ); + gate.open.store(true, Ordering::SeqCst); + driver.close().await.unwrap(); + driver.close().await.unwrap(); + assert_eq!(count(&events, "turn_finish"), 2, "{phase}"); + assert_eq!(count(&events, "turn_end"), 2, "{phase}"); + assert_eq!(gate.completed.load(Ordering::SeqCst), 1, "{phase}"); + let observed = observed.lock().unwrap(); + let terminals: Vec<&TurnResult> = observed + .iter() + .filter_map(|event| match event { + AgentEvent::TurnFinished(result) => Some(result), + _ => None, + }) + .collect(); + assert_eq!(terminals.len(), 1, "{phase}"); + let result = terminals[0]; + assert_eq!( + result.finish_reason, + FinishReason::Completed, + "existing final outcome was resumed, not replaced" + ); + assert_eq!( + text(&result.items[0]), + "reply:mfirst:mlast:ffirst:flast", + "{phase}" + ); + assert_eq!( + driver.snapshot().transcript.last(), + result.items.last(), + "{phase}" + ); + } +} From 9afba53989c5641c98dacae305c4606792d6eec2 Mon Sep 17 00:00:00 2001 From: daniel Date: Thu, 8 Oct 2026 10:10:36 +0100 Subject: [PATCH 2/3] refactor(loop): reduce mutator lifecycle to interception hooks Narrow the LoopMutator extension to three awaited hooks that intercept values before the loop consumes or commits them: on_session_start, on_model_request, and on_model_response. Remove the read-only lifecycle notifications (turn start/finish/end, model progress/error, tool progress/batch, session end), queued input admission and submit_input_async, explicit close, resumable terminal delivery, transcript rewrite notifications, and the Clone derive on LoopError. Read-only notifications belong in LoopObserver. This also removes the expect-based state machines and behaviour that differed depending on whether a mutator was registered. Hooks fail closed: an error or invalid edit aborts the operation without committing anything from it. LoopError::Cancelled from a hook finishes the turn as cancelled, matching mutate and model-stream cancellation. Hooks run inside the chat span, so gen_ai.input.messages and gen_ai.output.messages capture the final edited messages and hook latency counts toward inference time. MutationPoint::AfterTurnEnded gets a doc fix only: it has always run before the first inference of a new turn, after pending input is appended. --- book/src/ch16-compaction.md | 2 +- crates/agentkit-loop/README.md | 16 +- crates/agentkit-loop/src/hook_tests.rs | 458 +++++ crates/agentkit-loop/src/lib.rs | 1678 +++++-------------- crates/agentkit-loop/src/lifecycle_tests.rs | 1415 ---------------- 5 files changed, 843 insertions(+), 2726 deletions(-) create mode 100644 crates/agentkit-loop/src/hook_tests.rs delete mode 100644 crates/agentkit-loop/src/lifecycle_tests.rs diff --git a/book/src/ch16-compaction.md b/book/src/ch16-compaction.md index bfacc54..32e2599 100644 --- a/book/src/ch16-compaction.md +++ b/book/src/ch16-compaction.md @@ -261,7 +261,7 @@ This is why caching is configured separately from compaction in agentkit. Compac ## Loop integration -Compactors register as `LoopMutator`s. The loop runs transcript mutators before inference at `MutationPoint::AfterToolResult` (when continuing after a tool round) or `MutationPoint::AfterTurnEnded` (when dispatching a new logical turn). Despite its name, `AfterTurnEnded` is not a completion notification: newly buffered input has already been appended when it runs. The trigger decides which points are relevant. +Compactors register as `LoopMutator`s. The loop runs every registered mutator at each `MutationPoint` — `AfterToolResult` (between tool results and the next inference call) and `AfterTurnEnded` (before the first inference call of a new turn, after pending input is appended). The trigger decides which points are relevant. When a compactor fires: diff --git a/crates/agentkit-loop/README.md b/crates/agentkit-loop/README.md index d47a981..62eeb64 100644 --- a/crates/agentkit-loop/README.md +++ b/crates/agentkit-loop/README.md @@ -19,19 +19,15 @@ This crate provides: Use it as the central coordinator between model providers, tool execution, and application UI or control flow. -## Awaited lifecycle extensions +## Interception hooks -Register a `LoopMutator` with `AgentBuilder::mutator` to intercept lifecycle boundaries without driving or copying the loop implementation. Its lifecycle methods are asynchronous and default to no-ops, so existing transcript mutators remain compatible. Registered callbacks run in order and later callbacks see earlier payload edits. +Besides `mutate`, a `LoopMutator` registered with `AgentBuilder::mutator` can intercept values before the loop consumes or commits them: -The extension surface covers session startup and explicit closure, input admission, logical-turn startup and finalization, model requests and complete responses, model errors and stream progress, and resolved top-level tool batches. A logical turn may include multiple model requests and tool continuations; turn callbacks do not run once per inference. +- `on_session_start` edits session options before the model adapter starts the session. +- `on_model_request` edits a single inference request (transcript, tools, cache, metadata). These edits are not persisted. +- `on_model_response` edits complete model output before it is committed, returned, or sent to tools. Content may change. Item identity, accounting and tool-call linkage may not. -Payload types separate editable content from read-only identity and accounting. Complete response edits are applied before transcript commit and tool dispatch. Final-response edits are reflected in both the transcript and returned turn result. Stream-progress callbacks do not retract deltas already delivered to observers. - -Use `LoopDriver::submit_input_async` when input must pass awaited callbacks before acceptance. Synchronous submission defers that work until the next driver step. Use `LoopDriver::close().await` to finalize an owned session: a prompt finishing is not session closure, and dropping a driver cannot run asynchronous callbacks. - -Lifecycle extensions complement existing boundaries rather than replacing them: decorate a `ToolExecutor` for individual or nested tool interception, and a compaction implementation for an actual compaction operation. Model-selection and application-owned configuration changes remain outside the loop. - -The original `mutate` callback still runs at the two pre-inference transcript mutation points. `AfterTurnEnded` names new-turn dispatch, not a completion notification; `AfterToolResult` names a tool-round continuation. Use lifecycle callbacks for actual start/end notifications. +All hooks default to no-ops and run in registration order. Read-only notifications belong in `LoopObserver`; individual tool interception belongs at `ToolExecutor`. ## Quick start diff --git a/crates/agentkit-loop/src/hook_tests.rs b/crates/agentkit-loop/src/hook_tests.rs new file mode 100644 index 0000000..f0b703a --- /dev/null +++ b/crates/agentkit-loop/src/hook_tests.rs @@ -0,0 +1,458 @@ +use super::*; +use agentkit_tools_core::{ToolContext, ToolExecutionOutcome, ToolResult}; +use std::sync::{ + Mutex, + atomic::{AtomicBool, Ordering}, +}; + +#[derive(Clone, Copy)] +enum Script { + Reply, + Tool, +} + +#[derive(Clone, Default)] +struct ProbeModel { + scripts: Arc>>, + requests: Arc>>, + configs: Arc>>, +} + +impl ProbeModel { + fn new(scripts: impl IntoIterator) -> Self { + Self { + scripts: Arc::new(Mutex::new(scripts.into_iter().collect())), + ..Self::default() + } + } +} + +struct ProbeSession(ProbeModel); +struct ProbeTurn(VecDeque); + +#[async_trait] +impl ModelAdapter for ProbeModel { + type Session = ProbeSession; + async fn start_session(&self, config: SessionConfig) -> Result { + self.configs.lock().unwrap().push(config); + Ok(ProbeSession(self.clone())) + } +} + +#[async_trait] +impl ModelSession for ProbeSession { + type Turn = ProbeTurn; + async fn begin_turn( + &mut self, + request: TurnRequest, + _: Option, + ) -> Result { + self.0.requests.lock().unwrap().push(request); + let script = self.0.scripts.lock().unwrap().pop_front(); + let (item, finish_reason) = match script { + Some(Script::Reply) => ( + Item::text(ItemKind::Assistant, "reply").with_id("message-reply"), + FinishReason::Completed, + ), + Some(Script::Tool) => ( + Item::new( + ItemKind::Assistant, + vec![Part::ToolCall(ToolCallPart::new( + "call-probe", + "probe", + serde_json::json!({"value": 1}), + ))], + ) + .with_id("message-tool"), + FinishReason::ToolCall, + ), + None => return Err(LoopError::Provider("unexpected inference".into())), + }; + let mut events = VecDeque::new(); + if let Some(Part::ToolCall(call)) = item.parts.first() { + events.push_back(ModelTurnEvent::ToolCall(call.clone())); + } + events.push_back(ModelTurnEvent::Finished(ModelTurnResult { + finish_reason, + output_items: vec![item], + usage: None, + metadata: MetadataMap::new(), + model: None, + response_id: None, + })); + Ok(ProbeTurn(events)) + } +} + +#[async_trait] +impl ModelTurn for ProbeTurn { + async fn next_event( + &mut self, + _: Option, + ) -> Result, LoopError> { + Ok(self.0.pop_front()) + } +} + +#[derive(Clone, Default)] +struct ProbeExecutor(Arc>>); + +#[async_trait] +impl ToolExecutor for ProbeExecutor { + fn specs(&self) -> Vec { + Vec::new() + } + async fn execute(&self, request: ToolRequest, _: &mut ToolContext<'_>) -> ToolExecutionOutcome { + self.0.lock().unwrap().push(request.input); + ToolExecutionOutcome::Completed(ToolResult { + result: ToolResultPart { + call_id: request.call_id, + output: ToolOutput::Text("tool-result".into()), + is_error: false, + metadata: MetadataMap::new(), + }, + duration: None, + metadata: MetadataMap::new(), + }) + } +} + +/// Appends its label to every editable value so ordering is observable. +struct Tag(&'static str); + +#[async_trait] +impl LoopMutator for Tag { + async fn on_session_start( + &self, + payload: SessionStartPayload<'_>, + ctx: HookCtx<'_>, + ) -> Result<(), LoopError> { + assert!(ctx.turn_id.is_none()); + let seen = payload + .metadata + .get("order") + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_owned(); + payload + .metadata + .insert("order".into(), format!("{seen}{}", self.0).into()); + Ok(()) + } + + async fn on_model_request( + &self, + payload: ModelRequestPayload<'_>, + ctx: HookCtx<'_>, + ) -> Result<(), LoopError> { + assert!(ctx.turn_id.is_some()); + payload + .transcript + .push(Item::text(ItemKind::User, format!("ephemeral-{}", self.0))); + Ok(()) + } + + async fn on_model_response( + &self, + payload: ModelResponsePayload<'_>, + _: HookCtx<'_>, + ) -> Result<(), LoopError> { + for item in payload.output { + for part in item.parts.iter_mut() { + match part { + Part::Text(text) => text.text.push_str(self.0), + Part::ToolCall(call) => call.input["tag"] = self.0.into(), + _ => {} + } + } + } + Ok(()) + } +} + +struct ForgeToolCall; + +#[async_trait] +impl LoopMutator for ForgeToolCall { + async fn on_model_response( + &self, + payload: ModelResponsePayload<'_>, + _: HookCtx<'_>, + ) -> Result<(), LoopError> { + for item in payload.output { + for part in item.parts.iter_mut() { + if let Part::ToolCall(call) = part { + call.id = ToolCallId::new("forged"); + } + } + } + Ok(()) + } +} + +struct OrphanToolResult; + +#[async_trait] +impl LoopMutator for OrphanToolResult { + async fn on_model_request( + &self, + payload: ModelRequestPayload<'_>, + _: HookCtx<'_>, + ) -> Result<(), LoopError> { + payload.transcript.push(Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart { + call_id: ToolCallId::new("orphan"), + output: ToolOutput::Text("orphan".into()), + is_error: false, + metadata: MetadataMap::new(), + })], + )); + Ok(()) + } +} + +fn agent( + model: &ProbeModel, + executor: &ProbeExecutor, + mutators: impl FnOnce(AgentBuilder) -> AgentBuilder, +) -> Agent { + let builder = Agent::builder() + .model(model.clone()) + .tool_executor(executor.clone()) + .input(vec![Item::text(ItemKind::User, "hi")]); + mutators(builder).build().unwrap() +} + +fn text(item: &Item) -> &str { + match &item.parts[0] { + Part::Text(text) => &text.text, + other => panic!("expected text, got {other:?}"), + } +} + +#[tokio::test] +async fn session_start_edits_reach_the_adapter_in_registration_order() { + let model = ProbeModel::new([]); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder.mutator(Tag("a")).mutator(Tag("b")) + }); + agent.start(SessionConfig::new("session")).await.unwrap(); + let configs = model.configs.lock().unwrap(); + assert_eq!(configs[0].metadata["order"], "ab"); +} + +#[tokio::test] +async fn request_edits_are_inference_local_and_response_edits_are_committed() { + let model = ProbeModel::new([Script::Reply]); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder.mutator(Tag("a")).mutator(Tag("b")) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("expected finished turn"); + }; + + let request = model.requests.lock().unwrap()[0].clone(); + let sent: Vec<&str> = request.transcript.iter().map(text).collect(); + assert_eq!(sent, ["hi", "ephemeral-a", "ephemeral-b"]); + + assert_eq!(text(&result.items[0]), "replyab"); + let transcript = driver.snapshot().transcript; + let committed: Vec<&str> = transcript.iter().map(text).collect(); + assert_eq!(committed, ["hi", "replyab"]); +} + +#[tokio::test] +async fn response_edits_to_tool_arguments_reach_the_executor() { + let model = ProbeModel::new([Script::Tool, Script::Reply]); + let executor = ProbeExecutor::default(); + let agent = agent(&model, &executor, |builder| builder.mutator(Tag("a"))); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + loop { + match driver.next().await.unwrap() { + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => continue, + LoopStep::Finished(_) => break, + other => panic!("unexpected step {other:?}"), + } + } + assert_eq!( + executor.0.lock().unwrap().as_slice(), + [serde_json::json!({"value": 1, "tag": "a"})] + ); +} + +#[tokio::test] +async fn forged_tool_linkage_is_rejected_before_commit_or_execution() { + let model = ProbeModel::new([Script::Tool]); + let executor = ProbeExecutor::default(); + let agent = agent(&model, &executor, |builder| builder.mutator(ForgeToolCall)); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + let error = driver.next().await.unwrap_err(); + assert!(matches!(error, LoopError::Mutator(_)), "{error:?}"); + assert!(executor.0.lock().unwrap().is_empty()); + assert!( + driver + .snapshot() + .transcript + .iter() + .all(|item| item.kind != ItemKind::Assistant) + ); +} + +#[tokio::test] +async fn protocol_invalid_request_edits_are_rejected_before_inference() { + let model = ProbeModel::new([Script::Reply]); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder.mutator(OrphanToolResult) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + let error = driver.next().await.unwrap_err(); + assert!(matches!(error, LoopError::Mutator(_)), "{error:?}"); + assert!(model.requests.lock().unwrap().is_empty()); +} + +#[derive(Clone, Copy, PartialEq)] +enum Hook { + SessionStart, + Request, + Response, +} + +/// Fails `hook` once with `error`, then behaves as a no-op. +struct FailOnce { + hook: Hook, + error: fn() -> LoopError, + armed: AtomicBool, +} + +impl FailOnce { + fn new(hook: Hook, error: fn() -> LoopError) -> Self { + Self { + hook, + error, + armed: true.into(), + } + } + + fn check(&self, hook: Hook) -> Result<(), LoopError> { + if self.hook == hook && self.armed.swap(false, Ordering::SeqCst) { + Err((self.error)()) + } else { + Ok(()) + } + } +} + +#[async_trait] +impl LoopMutator for FailOnce { + async fn on_session_start( + &self, + _: SessionStartPayload<'_>, + _: HookCtx<'_>, + ) -> Result<(), LoopError> { + self.check(Hook::SessionStart) + } + + async fn on_model_request( + &self, + _: ModelRequestPayload<'_>, + _: HookCtx<'_>, + ) -> Result<(), LoopError> { + self.check(Hook::Request) + } + + async fn on_model_response( + &self, + _: ModelResponsePayload<'_>, + _: HookCtx<'_>, + ) -> Result<(), LoopError> { + self.check(Hook::Response) + } +} + +fn hook_error() -> LoopError { + LoopError::Mutator("hook failed".into()) +} + +#[tokio::test] +async fn session_start_failure_never_starts_the_adapter_session() { + let model = ProbeModel::new([]); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder.mutator(FailOnce::new(Hook::SessionStart, hook_error)) + }); + let Err(error) = agent.start(SessionConfig::new("session")).await else { + panic!("expected start to fail"); + }; + assert!(matches!(error, LoopError::Mutator(_)), "{error:?}"); + assert!(model.configs.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn response_failure_commits_nothing_and_the_session_recovers() { + let model = ProbeModel::new([Script::Tool, Script::Reply]); + let executor = ProbeExecutor::default(); + let agent = agent(&model, &executor, |builder| { + builder.mutator(FailOnce::new(Hook::Response, hook_error)) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + let error = driver.next().await.unwrap_err(); + assert!(matches!(error, LoopError::Mutator(_)), "{error:?}"); + assert!(executor.0.lock().unwrap().is_empty()); + let kinds: Vec = driver + .snapshot() + .transcript + .iter() + .map(|item| item.kind) + .collect(); + assert_eq!(kinds, [ItemKind::User]); + + let request = match driver.next().await.unwrap() { + LoopStep::Interrupt(LoopInterrupt::AwaitingInput(request)) => request, + other => panic!("unexpected step {other:?}"), + }; + request + .submit(&mut driver, vec![Item::text(ItemKind::User, "again")]) + .unwrap(); + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("expected finished turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Completed); +} + +#[tokio::test] +async fn cancelled_request_hook_finishes_cancelled_without_inference() { + let model = ProbeModel::new([Script::Reply]); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder.mutator(FailOnce::new(Hook::Request, || LoopError::Cancelled)) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("expected finished turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Cancelled); + assert!(model.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn cancelled_response_hook_finishes_cancelled_without_committing_output() { + let model = ProbeModel::new([Script::Tool]); + let executor = ProbeExecutor::default(); + let agent = agent(&model, &executor, |builder| { + builder.mutator(FailOnce::new(Hook::Response, || LoopError::Cancelled)) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("expected finished turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Cancelled); + assert!(executor.0.lock().unwrap().is_empty()); + assert!( + driver + .snapshot() + .transcript + .iter() + .flat_map(|item| &item.parts) + .all(|part| !matches!(part, Part::ToolCall(_))) + ); +} diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index 28fe115..4628613 100644 --- a/crates/agentkit-loop/src/lib.rs +++ b/crates/agentkit-loop/src/lib.rs @@ -792,9 +792,9 @@ pub struct ObservedEvent { /// replication, or audit. /// /// Observers are called *synchronously* from inside the driver, in the -/// same order items land in the transcript. Replacements are delivered through -/// [`TranscriptObserver::on_transcript_rewrite`], not as fictitious appends. -/// Persistence consumers that enable transcript editing must handle both methods. +/// same order items land in the transcript. Compaction-driven transcript +/// rewrites do **not** fire `on_transcript_event` — those are signaled by +/// [`AgentEvent::CompactionFinished`] instead. /// /// Register via [`AgentBuilder::transcript_observer`]; multiple observers /// may be registered and are called in registration order. @@ -820,10 +820,6 @@ pub trait TranscriptObserver: Send + Sync { /// state behind interior mutability so the driver can share an /// `Arc`. fn on_transcript_event(&self, event: TranscriptEvent<'_>); - - /// Replace the persisted transcript with this complete committed snapshot. - /// The default preserves compatibility for append-only observers. - fn on_transcript_rewrite(&self, _event: TranscriptRewriteEvent<'_>) {} } /// Session-addressed transcript append event delivered to @@ -836,13 +832,6 @@ pub struct TranscriptEvent<'a> { pub item: &'a Item, } -/// A complete committed replacement, including staged turn-start input. -#[derive(Clone, Debug)] -pub struct TranscriptRewriteEvent<'a> { - pub session_id: &'a SessionId, - pub items: &'a [Item], -} - /// Where in the loop a [`LoopMutator`] is given a chance to modify the /// transcript. Mutators run synchronously at these points; mid-stream /// mutation (e.g. between content deltas) is intentionally not supported @@ -853,9 +842,8 @@ pub enum MutationPoint { /// A tool result has just been appended; the next loop step will be /// another inference call. AfterToolResult, - /// Legacy name for the pre-inference transcript mutation pass on a new - /// logical turn, after input append. This is not a terminal callback; - /// use [`LoopMutator::on_turn_end`] for actual logical retirement. + /// A new logical turn is about to make its first inference call. Pending + /// user input has already been appended to the transcript. AfterTurnEnded, } @@ -907,93 +895,69 @@ impl<'a> std::ops::DerefMut for TranscriptCursor<'a> { } } -/// Read-only identity and cancellation context for awaited lifecycle callbacks. -/// -/// A logical turn can contain several model calls and approval/tool continuations. -/// `model_call_index` is session-local and counts calls, not provider retry attempts. -/// Terminal callbacks run even when the cancellation token is cancelled. +/// Read-only context handed to the [`LoopMutator`] interception hooks. #[non_exhaustive] -pub struct LifecycleCtx<'a> { +pub struct HookCtx<'a> { + /// Session the hook runs in. pub session_id: &'a SessionId, + /// Turn the hook is associated with; `None` at session start. pub turn_id: Option<&'a agentkit_core::TurnId>, - pub model_call_index: Option, + /// Cancellation handle for the active turn, if any. pub cancellation: Option, - /// Selection snapshot; per-call routing inside `begin_turn` can change it. - pub provider_name: Option<&'a str>, - pub model_name: Option<&'a str>, - /// Original failure at a terminal boundary, never replaced by hook errors. - pub error: Option<&'a LoopError>, } -/// Read-only state at explicit session closure. No detached work is awaited. -pub struct SessionEndPayload<'a> { - pub transcript: &'a [Item], -} - -/// Loop-owned logical progress. This is not arbitrary progress reported by tools. -pub enum TurnProgress<'a> { - Model(&'a ModelTurnEvent), - ToolDetached(&'a ToolResultPart), -} - -/// Editable session options; the session identity is available only in the context. +/// Editable session options, applied before the model adapter consumes them. +#[non_exhaustive] pub struct SessionStartPayload<'a> { + /// Metadata forwarded to the model adapter. pub metadata: &'a mut MetadataMap, + /// Default prompt caching policy for the session. pub cache: &'a mut Option, + /// Features the consumer of model-turn events can handle. pub consumer_capabilities: &'a mut SessionConsumerCapabilities, } -/// Inference-local request options. Transcript edits affect only this model call; -/// persistent history edits use [`LoopMutator::mutate`] or `on_turn_start`. -/// Catalog changes do not install executable tools. Session/turn identity is read-only. -/// Prompt edits are validated before the request is consumed. +/// Editable inference request. Edits apply to this model call only; persistent +/// transcript edits belong in [`LoopMutator::mutate`]. Editing +/// `available_tools` changes what the model sees, not what the executor runs. +#[non_exhaustive] pub struct ModelRequestPayload<'a> { + /// Transcript sent to the model. Tool-call pairing is validated after hooks. pub transcript: &'a mut Vec, + /// Tool specs advertised to the model. pub available_tools: &'a mut Vec, + /// Prompt caching policy for this call. pub cache: &'a mut Option, + /// Request metadata forwarded to the model adapter. pub metadata: &'a mut MetadataMap, } -/// Editable content of one output item. Identity, role, accounting and timestamp -/// remain read-only in `original`, the snapshot at entry to this callback. -/// `parts` and `metadata` expose the current candidate, including earlier hooks. -/// Tool linkage in `parts` is checked before commit. +/// Editable content of one output item. `original` is the item as it entered +/// this hook, including edits from earlier mutators. +#[non_exhaustive] pub struct OutputItemPayload<'a> { + /// The item as it entered this hook. pub original: &'a Item, + /// Editable content parts. pub parts: &'a mut Vec, + /// Editable item metadata. pub metadata: &'a mut MetadataMap, } -/// Complete, uncommitted model output. Item identities, roles, accounting, timestamps, -/// and tool-call/result identities must remain unchanged; content can be rewritten. -/// The driver checks these constraints before committing any output. -/// Already emitted streaming deltas cannot be retracted by this callback. +/// Complete model output before it is committed to the transcript, returned +/// in the turn result, or dispatched to tools. Content may be rewritten; item +/// identity, roles, accounting and tool-call linkage (ids and names) may not. +/// Deltas already streamed to observers are not retracted. +#[non_exhaustive] pub struct ModelResponsePayload<'a> { + /// One entry per output item, in order. pub output: Vec>, + /// Why the model stopped generating. pub finish_reason: &'a FinishReason, + /// Token usage for the call, if reported. pub usage: Option<&'a Usage>, + /// Provider metadata for the call. pub metadata: &'a MetadataMap, - pub model: Option<&'a str>, - pub response_id: Option<&'a str>, -} - -/// Logical finalization candidate. `result` is the immutable snapshot at entry -/// to this callback, including earlier hooks' edits. -/// `output` is editable only for uncommitted successful final assistant output. -/// Cancellation, failure, and already committed tool-bearing output are read-only. -/// Output edits obey the same constraints as [`ModelResponsePayload`]. -pub struct TurnFinishPayload<'a> { - pub result: &'a TurnResult, - pub output: Option>>, -} - -/// A transcript-resolved top-level tool batch, not a nested execution batch. -/// Results include cancellation/denial and background-detachment placeholders; -/// a placeholder does not imply that background execution has completed. -pub struct ToolBatchPayload<'a> { - /// Model-issued calls. Approved argument rewrites belong to executor facts. - pub calls: &'a [ToolCallPart], - pub results: &'a [ToolResultPart], } /// Async transcript mutator. Registered via [`AgentBuilder::mutator`] and @@ -1001,136 +965,46 @@ pub struct ToolBatchPayload<'a> { /// (e.g. running token totals via interior mutability) and decide for /// themselves whether and how to modify the transcript. /// -/// Every method defaults to a no-op. Lifecycle callbacks are awaited in -/// registration order at their typed boundaries; `mutate` remains the generic -/// transcript editing seam. Staged edits are committed only after validation. -/// Tool/executor, compaction, routing and application configuration remain -/// separate extension seams; this trait does not intercept nested tools. +/// The `on_*` hooks intercept values before the loop consumes or commits +/// them. Every method defaults to a no-op, and registered mutators run in +/// order, each seeing earlier mutators' edits. Read-only notifications belong +/// in [`LoopObserver`]; individual tool interception belongs at +/// [`ToolExecutor`]. +/// +/// Hooks fail closed: an error or an invalid edit aborts the operation and +/// commits nothing from it, so an edit that did not apply is never silently +/// skipped. [`Agent::start`] returns the error. Inside a turn, the turn +/// finishes with an error, or as cancelled when the hook returns +/// [`LoopError::Cancelled`]. #[async_trait] pub trait LoopMutator: Send + Sync { - /// Runs before session options are consumed by the model adapter. + /// Edit session options before the model adapter starts the session. async fn on_session_start( &self, - _payload: SessionStartPayload<'_>, - _ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - Ok(()) - } - - /// Runs once on explicit [`LoopDriver::close`], not on turn completion or Drop. - async fn on_session_end( - &self, - _payload: SessionEndPayload<'_>, - _ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - Ok(()) - } - - /// Runs before input dispatch. Sync submission queues admission until `next`; - /// [`LoopDriver::submit_input_async`] awaits this callback before acceptance. - async fn on_input( - &self, - _input: &mut Vec, - _ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - Ok(()) - } - - /// Runs once before work in a newly active logical turn. The cursor is a - /// staged transcript including admitted input. Edits are validated before - /// commit; failure leaves the live transcript unchanged and retires the turn. - async fn on_turn_start( - &self, - _transcript: &mut TranscriptCursor<'_>, - _ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - Ok(()) - } - - /// Runs before logical finalization, including cancellation/error (read-only). - async fn on_turn_finish( - &self, - _payload: TurnFinishPayload<'_>, - _ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - Ok(()) - } - - /// Runs once after logical retirement. Every registered callback is attempted. - /// On ordinary retirement, errors are logged without replacing the committed - /// result. Explicit `close` aggregates notification errors in its close result. - async fn on_turn_end( - &self, - _result: &TurnResult, - _ctx: LifecycleCtx<'_>, + payload: SessionStartPayload<'_>, + ctx: HookCtx<'_>, ) -> Result<(), LoopError> { + let _ = (payload, ctx); Ok(()) } - /// Runs immediately before the inference request is consumed by the session. + /// Edit an inference request before the model session consumes it. async fn on_model_request( &self, - _payload: ModelRequestPayload<'_>, - _ctx: LifecycleCtx<'_>, + payload: ModelRequestPayload<'_>, + ctx: HookCtx<'_>, ) -> Result<(), LoopError> { + let _ = (payload, ctx); Ok(()) } - /// Runs on complete output, before transcript commit and tool dispatch. + /// Edit complete model output before commit and tool dispatch. async fn on_model_response( &self, - _payload: ModelResponsePayload<'_>, - _ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - Ok(()) - } - - /// Read-only model failure notification, including cancellation. Hook errors - /// never replace the model error. This does not report tool or hook failures. - async fn on_model_error( - &self, - _error: &LoopError, - _ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - Ok(()) - } - - /// Read-only logical progress before forwarding model events or committing - /// a background-detachment placeholder. Individual tool progress is external. - async fn on_turn_progress( - &self, - _progress: TurnProgress<'_>, - _ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - Ok(()) - } - - /// Read-only notification of loop-owned background-detachment progress. - async fn on_tool_progress( - &self, - _progress: &ToolResultPart, - _ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - Ok(()) - } - - /// Read-only stream/progress notification, before observer forwarding. - /// Includes usage, tool requests, Finished, and response-attempt supersession. - async fn on_model_progress( - &self, - _event: &ModelTurnEvent, - _ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - Ok(()) - } - - /// Runs once when every top-level call in a model batch has a transcript result. - /// Individual and nested tool interception belongs at [`ToolExecutor`]. - async fn on_tool_batch( - &self, - _payload: ToolBatchPayload<'_>, - _ctx: LifecycleCtx<'_>, + payload: ModelResponsePayload<'_>, + ctx: HookCtx<'_>, ) -> Result<(), LoopError> { + let _ = (payload, ctx); Ok(()) } @@ -1543,47 +1417,9 @@ struct ActiveToolRound { foreground_progressed: bool, } -struct PendingToolBatch { - turn_id: agentkit_core::TurnId, - calls: Vec, - results: BTreeMap, - next_mutator: usize, - error: Option, -} - -struct TerminalState { - result: TurnResult, - original_items: Vec, - editable: bool, - output_committed: bool, - finish_index: usize, - finish_done: bool, - terminal_emitted: bool, - end_index: usize, - failure: Option, - finish_error: Option, - end_error: Option, -} - -struct ClosingState { - cleanup_turns: Vec, - cleanup_index: usize, - cleanup_done: bool, - session_index: usize, - error: Option, -} - #[derive(Default)] struct DriverLifecycle { active_turn: Option, - closed: bool, - closing: Option, - terminal: Option, - model_call_index: u64, - pending_cancel: Option, - finish_attempted: bool, - cancellation: Option, - started_with_input: bool, } /// A configured agent ready to start a session. @@ -1676,43 +1512,17 @@ where ) -> Result, LoopError> { let session_id = config.session_id.clone(); for mutator in &self.mutators { - mutator - .on_session_start( - SessionStartPayload { - metadata: &mut config.metadata, - cache: &mut config.cache, - consumer_capabilities: &mut config.consumer_capabilities, - }, - LifecycleCtx { - session_id: &session_id, - turn_id: None, - model_call_index: None, - cancellation: None, - provider_name: self.model.provider_name(), - model_name: None, - error: None, - }, - ) - .await?; - } - let mut input = self.input.clone(); - if !input.is_empty() { - for mutator in &self.mutators { - mutator - .on_input( - &mut input, - LifecycleCtx { - session_id: &session_id, - turn_id: None, - model_call_index: None, - cancellation: None, - provider_name: self.model.provider_name(), - model_name: None, - error: None, - }, - ) - .await?; - } + let payload = SessionStartPayload { + metadata: &mut config.metadata, + cache: &mut config.cache, + consumer_capabilities: &mut config.consumer_capabilities, + }; + let ctx = HookCtx { + session_id: &session_id, + turn_id: None, + cancellation: None, + }; + mutator.on_session_start(payload, ctx).await?; } let default_cache = config.cache.clone(); let mut session = self.model.start_session(config).await?; @@ -1749,9 +1559,7 @@ where observers: self.observers.clone(), transcript_observers: self.transcript_observers.clone(), transcript: self.transcript.clone(), - pending_input: input, - pending_admission: VecDeque::new(), - pending_tool_batches: Vec::new(), + pending_input: self.input.clone(), pending_approvals: BTreeMap::new(), pending_approval_order: VecDeque::new(), active_tool_round: None, @@ -1765,12 +1573,6 @@ where tool_cancellations: HashMap::new(), }; driver.emit(AgentEvent::RunStarted { session_id }); - if !self.mutators.is_empty() && !driver.pending_input.is_empty() { - driver.emit(AgentEvent::InputAccepted { - session_id: driver.session_id.clone(), - items: driver.pending_input.clone(), - }); - } Ok(driver) } } @@ -2033,8 +1835,6 @@ where transcript_observers: Vec>, transcript: Vec, pending_input: Vec, - pending_admission: VecDeque>, - pending_tool_batches: Vec, pending_approvals: BTreeMap, pending_approval_order: VecDeque, active_tool_round: Option, @@ -2163,332 +1963,34 @@ where !self.pending_approvals.is_empty() } - async fn start_logical_turn(&mut self) -> Result { + fn start_logical_turn(&mut self) -> agentkit_core::TurnId { if let Some(turn_id) = &self.lifecycle.active_turn { - return Ok(turn_id.clone()); + return turn_id.clone(); } let turn_id = agentkit_core::TurnId::new(format!("turn-{}", self.next_turn_index)); self.next_turn_index += 1; - self.start_logical_turn_with(turn_id).await + self.start_logical_turn_with(turn_id) } - async fn start_logical_turn_with( - &mut self, - turn_id: agentkit_core::TurnId, - ) -> Result { - if self.lifecycle.active_turn.as_ref() == Some(&turn_id) { - return Ok(turn_id); + fn start_logical_turn_with(&mut self, turn_id: agentkit_core::TurnId) -> agentkit_core::TurnId { + if let Some(active_turn) = &self.lifecycle.active_turn { + return active_turn.clone(); } self.lifecycle.active_turn = Some(turn_id.clone()); self.emit(AgentEvent::TurnStarted { session_id: self.session_id.clone(), turn_id: turn_id.clone(), }); - self.lifecycle.finish_attempted = false; - self.lifecycle.cancellation = self - .cancellation - .as_ref() - .map(CancellationHandle::checkpoint); - if !self.mutators.is_empty() { - let mut candidate = self.transcript.clone(); - self.lifecycle.started_with_input = !self.pending_input.is_empty(); - candidate.extend(std::mem::take(&mut self.pending_input)); - let original = candidate.clone(); - let mut cursor = TranscriptCursor { - items: &mut candidate, - dirty: false, - }; - for mutator in &self.mutators { - mutator - .on_turn_start(&mut cursor, self.lifecycle_ctx(Some(&turn_id), None)) - .await?; - } - validate_prompt_edits(&original, &candidate)?; - self.commit_transcript(candidate); - } - Ok(turn_id) + turn_id } - async fn finish_logical_turn( - &mut self, - result: &TurnResult, - failure: Option<&LoopError>, - ) -> Result<(), LoopError> { + fn finish_logical_turn(&mut self, result: &TurnResult) { if self.pending_round_resume.as_ref() == Some(&result.turn_id) { self.pending_round_resume = None; } - if self.lifecycle.terminal.is_none() { - self.begin_terminal(result, false, failure, false); - self.lifecycle - .terminal - .as_mut() - .expect("terminal state") - .finish_done = true; - } - let state = self.lifecycle.terminal.as_mut().expect("terminal state"); - state.result = result.clone(); - state.output_committed = true; - if let Some(error) = failure { - state.failure = Some(error.clone()); - } - if !state.terminal_emitted { - if self.lifecycle.active_turn.as_ref() == Some(&result.turn_id) { - self.emit(AgentEvent::TurnFinished(result.clone())); - } + if self.lifecycle.active_turn.as_ref() == Some(&result.turn_id) { self.lifecycle.active_turn = None; - self.lifecycle.started_with_input = false; - self.lifecycle - .terminal - .as_mut() - .expect("terminal state") - .terminal_emitted = true; - } - loop { - let state = self.lifecycle.terminal.as_ref().expect("terminal state"); - if state.end_index == self.mutators.len() { - break; - } - let outcome = self.mutators[state.end_index] - .on_turn_end( - &state.result, - LifecycleCtx { - error: state.failure.as_ref(), - ..self.lifecycle_ctx(Some(&state.result.turn_id), None) - }, - ) - .await; - let state = self.lifecycle.terminal.as_mut().expect("terminal state"); - state.end_index += 1; - if let Err(error) = outcome { - state.end_error.get_or_insert(error); - } - } - let state = self.lifecycle.terminal.take().expect("terminal state"); - if self.lifecycle.closing.is_some() { - return state.end_error.map_or(Ok(()), Err); - } - if let Some(error) = state.end_error { - tracing::debug!(%error, "turn end callback failed after retirement"); - } - Ok(()) - } - - fn begin_terminal( - &mut self, - result: &TurnResult, - editable: bool, - failure: Option<&LoopError>, - append_output: bool, - ) { - self.lifecycle.terminal = Some(TerminalState { - result: result.clone(), - original_items: result.items.clone(), - editable, - output_committed: !append_output, - finish_index: 0, - finish_done: false, - terminal_emitted: false, - end_index: 0, - failure: failure.cloned(), - finish_error: None, - end_error: None, - }); - } - - fn mark_terminal_output_committed(&mut self, turn_id: &agentkit_core::TurnId) { - if let Some(state) = &mut self.lifecycle.terminal { - if &state.result.turn_id == turn_id { - state.output_committed = true; - } - } - } - - fn lifecycle_ctx<'a>( - &'a self, - turn_id: Option<&'a agentkit_core::TurnId>, - model_call_index: Option, - ) -> LifecycleCtx<'a> { - LifecycleCtx { - session_id: &self.session_id, - turn_id, - model_call_index, - cancellation: if turn_id.is_some() { - self.lifecycle.cancellation.clone() - } else { - None - }, - provider_name: self - .session - .as_ref() - .and_then(ModelSession::provider_name) - .or(self.provider_name.as_deref()), - model_name: self.session.as_ref().and_then(ModelSession::model_name), - error: None, - } - } - - async fn run_turn_finish( - &mut self, - result: &mut TurnResult, - editable: bool, - failure: Option<&LoopError>, - append_output: bool, - ) -> Result<(), LoopError> { - self.lifecycle.finish_attempted = true; - if self.lifecycle.terminal.is_none() { - self.begin_terminal(result, editable, failure, append_output); - } - loop { - let state = self.lifecycle.terminal.as_ref().expect("terminal state"); - if state.finish_done || state.finish_index == self.mutators.len() { - break; - } - let visible = state.result.clone(); - let mut candidate = visible.clone(); - let output = if state.editable { - Some(output_payload(&visible.items, &mut candidate.items)) - } else { - None - }; - let outcome = self.mutators[state.finish_index] - .on_turn_finish( - TurnFinishPayload { - result: &visible, - output, - }, - LifecycleCtx { - error: state.failure.as_ref(), - ..self.lifecycle_ctx(Some(&visible.turn_id), None) - }, - ) - .await; - let state = self.lifecycle.terminal.as_mut().expect("terminal state"); - state.result = candidate; - state.finish_index += 1; - if let Err(error) = outcome { - state.finish_error.get_or_insert(error); - } - } - let state = self.lifecycle.terminal.as_mut().expect("terminal state"); - if !state.finish_done { - if let Err(error) = validate_output_edits(&state.original_items, &state.result.items) { - state.finish_error.get_or_insert(error); - } - state.finish_done = true; - } - *result = state.result.clone(); - state.finish_error.clone().map_or(Ok(()), Err) - } - - async fn resume_terminal_delivery(&mut self) -> Result { - let state = self.lifecycle.terminal.as_ref().expect("terminal state"); - let mut result = state.result.clone(); - let failure = state.failure.clone(); - let editable = state.editable; - let append_output = !state.output_committed; - let mut first_error = failure.clone(); - if let Err(error) = self - .run_turn_finish(&mut result, editable, failure.as_ref(), append_output) - .await - { - first_error.get_or_insert(error.clone()); - let state = self.lifecycle.terminal.as_mut().expect("terminal state"); - if !state.output_committed { - result.items.clear(); - result.finish_reason = FinishReason::Error; - state.result = result.clone(); - state.output_committed = true; - } - state.failure.get_or_insert(error); - } - if !self - .lifecycle - .terminal - .as_ref() - .expect("terminal state") - .output_committed - { - self.extend_transcript(result.items.clone()); - self.mark_terminal_output_committed(&result.turn_id); - self.close_interrupted_tool_calls(); - } - if let Err(error) = self.run_ready_tool_batches().await { - first_error.get_or_insert(error); - } - let failure = self - .lifecycle - .terminal - .as_ref() - .expect("terminal state") - .failure - .clone(); - if let Err(error) = self.finish_logical_turn(&result, failure.as_ref()).await { - first_error.get_or_insert(error); - } - first_error.map_or(Ok(result), Err) - } - - async fn notify_model_error(&mut self, error: &LoopError, turn_id: &agentkit_core::TurnId) { - for mutator in &self.mutators { - if let Err(hook_error) = mutator - .on_model_error( - error, - self.lifecycle_ctx(Some(turn_id), Some(self.lifecycle.model_call_index)), - ) - .await - { - tracing::debug!(%hook_error, "model error callback failed"); - } - } - } - - async fn run_ready_tool_batches(&mut self) -> Result<(), LoopError> { - let mut index = 0; - let mut first_error = None; - while index < self.pending_tool_batches.len() { - let batch = &self.pending_tool_batches[index]; - if batch.results.len() != batch.calls.len() { - index += 1; - continue; - } - while self.pending_tool_batches[index].next_mutator < self.mutators.len() { - let batch = &self.pending_tool_batches[index]; - let results: Vec = batch - .calls - .iter() - .map(|call| batch.results[&call.id].clone()) - .collect(); - let outcome = self.mutators[batch.next_mutator] - .on_tool_batch( - ToolBatchPayload { - calls: &batch.calls, - results: &results, - }, - self.lifecycle_ctx(Some(&batch.turn_id), None), - ) - .await; - let batch = &mut self.pending_tool_batches[index]; - batch.next_mutator += 1; - if let Err(error) = outcome { - batch.error.get_or_insert(error); - } - } - let batch = self.pending_tool_batches.remove(index); - if let Some(error) = batch.error { - first_error.get_or_insert(error); - } - } - first_error.map_or(Ok(()), Err) - } - - fn record_batch_result(&mut self, result: &ToolResultPart) { - for batch in &mut self.pending_tool_batches { - if batch.calls.iter().any(|call| call.id == result.call_id) { - batch - .results - .entry(result.call_id.clone()) - .or_insert_with(|| result.clone()); - } + self.emit(AgentEvent::TurnFinished(result.clone())); } } @@ -2603,7 +2105,7 @@ where self.collect_pending_loop_updates().await?; let mut resolutions = std::mem::take(&mut self.pending_loop_updates); if !resolutions.is_empty() { - self.start_logical_turn().await?; + self.start_logical_turn(); } let mut saw_items = false; while let Some(resolution) = resolutions.pop_front() { @@ -2613,7 +2115,7 @@ where saw_items = true; } TaskResolution::Approval(task) => { - let turn_id = self.start_logical_turn().await?; + let turn_id = self.start_logical_turn(); self.enqueue_pending_approval(&turn_id, task, None); } } @@ -2662,9 +2164,8 @@ where session_id: &observed_session_id, observers: &observers, }; - let mut candidate = self.transcript.clone(); let mut cursor = TranscriptCursor { - items: &mut candidate, + items: &mut self.transcript, dirty: false, }; for mutator in &mutators { @@ -2685,7 +2186,6 @@ where } if cursor.dirty { validate_transcript_invariants(cursor.items)?; - self.commit_transcript(candidate); } Ok(()) } @@ -2714,7 +2214,6 @@ where self.active_tool_round = None; return self .finish_cancelled(presentation_turn_id, Vec::new()) - .await .map(Some); } @@ -2765,8 +2264,7 @@ where TaskStartOutcome::Pending { kind, .. } => { self.emit(AgentEvent::ToolExecutionStarted(call.clone())); if kind == agentkit_task_manager::TaskKind::Background { - self.append_detach_placeholder(call.id.clone(), &call.name) - .await?; + self.append_detach_placeholder(call.id.clone(), &call.name); if let Some(active) = self.active_tool_round.as_mut() { active.background_pending = true; } @@ -2801,8 +2299,7 @@ where } } Some(TurnTaskUpdate::Detached(snapshot)) => { - self.append_detach_placeholder(snapshot.call_id, &snapshot.tool_name) - .await?; + self.append_detach_placeholder(snapshot.call_id, &snapshot.tool_name); if let Some(active) = self.active_tool_round.as_mut() { active.background_pending = true; active.foreground_progressed = true; @@ -2822,7 +2319,6 @@ where self.active_tool_round = None; return self .finish_cancelled(presentation_turn_id, Vec::new()) - .await .map(Some); } let active = self.active_tool_round.take().ok_or_else(|| { @@ -2877,33 +2373,17 @@ where turn_id: agentkit_core::TurnId, mutation_point: MutationPoint, ) -> Result { - // Preserve cancellation that arrived while awaiting a logical-start hook. - // Existing approval/tool continuations retain their checkpoint policy. - if mutation_point == MutationPoint::AfterTurnEnded - && self - .lifecycle - .cancellation - .as_ref() - .is_some_and(TurnCancellation::is_cancelled) - { - return self - .finish_cancelled(turn_id, interrupted_assistant_items()) - .await; - } let cancellation = self .cancellation .as_ref() .map(CancellationHandle::checkpoint); - self.lifecycle.cancellation = cancellation.clone(); match self .run_mutators(mutation_point, Some(&turn_id), cancellation.clone()) .await { Ok(()) => {} Err(LoopError::Cancelled) => { - return self - .finish_cancelled(turn_id, interrupted_assistant_items()) - .await; + return self.finish_cancelled(turn_id, interrupted_assistant_items()); } Err(error) => return Err(error), } @@ -2914,16 +2394,14 @@ where // for the model to respond to. Finish the turn rather than dispatch an // assistant-prefill request, which most providers reject. if !transcript_has_pending_input(&self.transcript) { - let mut turn_result = TurnResult { + let turn_result = TurnResult { turn_id, finish_reason: FinishReason::Completed, items: Vec::new(), usage: None, metadata: MetadataMap::new(), }; - self.run_turn_finish(&mut turn_result, true, None, true) - .await?; - self.finish_logical_turn(&turn_result, None).await?; + self.finish_logical_turn(&turn_result); return Ok(LoopStep::Finished(turn_result)); } @@ -2931,9 +2409,7 @@ where .as_ref() .is_some_and(TurnCancellation::is_cancelled) { - return self - .finish_cancelled(turn_id, interrupted_assistant_items()) - .await; + return self.finish_cancelled(turn_id, interrupted_assistant_items()); } let catalog_events = self.tool_executor.drain_catalog_events(); @@ -2950,35 +2426,17 @@ where .or_else(|| self.default_cache.clone()), metadata: MetadataMap::new(), }; - - self.run_ready_tool_batches().await?; - self.lifecycle.model_call_index += 1; - let original_request = request.transcript.clone(); - for mutator in &self.mutators { - mutator - .on_model_request( - ModelRequestPayload { - transcript: &mut request.transcript, - available_tools: &mut request.available_tools, - cache: &mut request.cache, - metadata: &mut request.metadata, - }, - self.lifecycle_ctx(Some(&turn_id), Some(self.lifecycle.model_call_index)), - ) - .await?; - } - if !self.mutators.is_empty() { - validate_prompt_edits(&original_request, &request.transcript)?; - } let session = self .session .as_mut() .ok_or_else(|| LoopError::InvalidState("model session is not available".into()))?; // Inference span per the OTel GenAI semantic conventions. It wraps the - // model request and the full event drain rather than just `begin_turn`, - // so attributes that streaming adapters only learn mid-stream (usage, - // stop reason, response identity) still land before the span closes. + // mutator request/response hooks, the model request, and the full event + // drain rather than just `begin_turn`, so recorded messages are the final + // edited ones, hook latency counts toward inference time, and attributes + // that streaming adapters only learn mid-stream (usage, stop reason, + // response identity) still land before the span closes. // `otel.name` carries the dynamic `chat {model}` span name for // OpenTelemetry bridges since tracing span names are static. let chat_span = tracing::info_span!( @@ -2998,14 +2456,6 @@ where "gen_ai.usage.output_tokens" = tracing::field::Empty, "gen_ai.usage.cost" = tracing::field::Empty, ); - if let Some(capture) = self.telemetry.input_messages { - record_string_array_attribute( - &chat_span, - "gen_ai.input.messages", - capture_messages(&request.transcript, capture, CaptureOrder::NewestTail), - ); - } - // Seed known identity before begin_turn so request setup failures and // cancellation remain attributable. Successful per-turn routing below // overwrites these values with the effective selection. @@ -3026,6 +2476,35 @@ where } use tracing::Instrument; + match run_request_hooks( + &self.mutators, + &self.session_id, + &turn_id, + cancellation.as_ref(), + &mut request, + ) + .instrument(chat_span.clone()) + .await + { + Ok(()) => {} + Err(LoopError::Cancelled) => { + self.task_manager + .on_turn_interrupted(&turn_id) + .await + .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; + return self.finish_cancelled(turn_id, interrupted_assistant_items()); + } + Err(error) => return Err(error), + } + + if let Some(capture) = self.telemetry.input_messages { + record_string_array_attribute( + &chat_span, + "gen_ai.input.messages", + capture_messages(&request.transcript, capture, CaptureOrder::NewestTail), + ); + } + let mut turn = match session .begin_turn(request, cancellation.clone()) .instrument(chat_span.clone()) @@ -3033,20 +2512,13 @@ where { Ok(turn) => turn, Err(LoopError::Cancelled) => { - self.notify_model_error(&LoopError::Cancelled, &turn_id) - .await; self.task_manager .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self - .finish_cancelled(turn_id, interrupted_assistant_items()) - .await; - } - Err(error) => { - self.notify_model_error(&error, &turn_id).await; - return Err(error); + return self.finish_cancelled(turn_id, interrupted_assistant_items()); } + Err(error) => return Err(error), }; // begin_turn may apply per-turn routing. Sample the effective selection @@ -3067,6 +2539,7 @@ where } } + let mut saw_tool_call = false; let mut finished_result = None; let mut latest_usage = None; let mut streamed_content = cancellation @@ -3080,39 +2553,20 @@ where { Ok(event) => event, Err(LoopError::Cancelled) => { - self.notify_model_error(&LoopError::Cancelled, &turn_id) - .await; self.task_manager .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self - .finish_cancelled(turn_id, interrupted_stream_items(streamed_content.as_ref())) - .await; - } - Err(error) => { - self.notify_model_error(&error, &turn_id).await; - return Err(error); + return self.finish_cancelled( + turn_id, + interrupted_stream_items(streamed_content.as_ref()), + ); } + Err(error) => return Err(error), } { - for mutator in &self.mutators { - mutator - .on_model_progress( - &event, - self.lifecycle_ctx(Some(&turn_id), Some(self.lifecycle.model_call_index)), - ) - .await?; - } - for mutator in &self.mutators { - mutator - .on_turn_progress( - TurnProgress::Model(&event), - self.lifecycle_ctx(Some(&turn_id), Some(self.lifecycle.model_call_index)), - ) - .await?; - } let attempt_superseded = matches!(event, ModelTurnEvent::ResponseAttemptSuperseded); if attempt_superseded { + saw_tool_call = false; latest_usage = None; if let Some(content) = &mut streamed_content { content.reset(); @@ -3123,16 +2577,15 @@ where .as_ref() .is_some_and(TurnCancellation::is_cancelled) { - self.notify_model_error(&LoopError::Cancelled, &turn_id) - .await; turn.on_cancelled(); self.task_manager .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self - .finish_cancelled(turn_id, interrupted_stream_items(streamed_content.as_ref())) - .await; + return self.finish_cancelled( + turn_id, + interrupted_stream_items(streamed_content.as_ref()), + ); } if attempt_superseded { continue; @@ -3149,6 +2602,7 @@ where self.emit(AgentEvent::UsageUpdated(usage)); } ModelTurnEvent::ToolCall(call) => { + saw_tool_call = true; if let Some(content) = &mut streamed_content { content.commit_tool_call(&call); } @@ -3164,31 +2618,33 @@ where } } - let Some(mut result) = finished_result else { - let error = LoopError::Provider("model turn ended without a Finished event".into()); - self.notify_model_error(&error, &turn_id).await; - return Err(error); - }; + let mut result = finished_result.ok_or_else(|| { + LoopError::Provider("model turn ended without a Finished event".into()) + })?; result.usage = merge_usage(result.usage, latest_usage); - let original_output = result.output_items.clone(); - for mutator in &self.mutators { - let visible = result.output_items.clone(); - mutator - .on_model_response( - ModelResponsePayload { - output: output_payload(&visible, &mut result.output_items), - finish_reason: &result.finish_reason, - usage: result.usage.as_ref(), - metadata: &result.metadata, - model: result.model.as_deref(), - response_id: result.response_id.as_deref(), - }, - self.lifecycle_ctx(Some(&turn_id), Some(self.lifecycle.model_call_index)), - ) - .await?; + match run_response_hooks( + &self.mutators, + &self.session_id, + &turn_id, + cancellation.as_ref(), + &mut result, + ) + .instrument(chat_span.clone()) + .await + { + Ok(()) => {} + Err(LoopError::Cancelled) => { + self.task_manager + .on_turn_interrupted(&turn_id) + .await + .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; + return self.finish_cancelled( + turn_id, + interrupted_stream_items(streamed_content.as_ref()), + ); + } + Err(error) => return Err(error), } - validate_output_edits(&original_output, &result.output_items)?; - let saw_tool_call = !extract_tool_calls(&result.output_items).is_empty(); if let Some(model) = &result.model { chat_span.record("gen_ai.response.model", model.as_str()); } @@ -3215,7 +2671,14 @@ where "gen_ai.response.finish_reasons", provider_finish_reasons(&result.metadata, &result.finish_reason), ); - + if let Some(capture) = self.telemetry.output_messages { + record_string_array_attribute( + &chat_span, + "gen_ai.output.messages", + capture_messages(&result.output_items, capture, CaptureOrder::OldestHead), + ); + } + drop(chat_span); tracing::Span::current().record("saw_tool_call", saw_tool_call); tracing::Span::current().record( "finish_reason", @@ -3231,52 +2694,20 @@ where if matches!(item.kind, ItemKind::Assistant) { if item.usage.is_none() { item.usage = usage.clone(); - } - if item.finish_reason.is_none() { - item.finish_reason = Some(finish_reason.clone()); - } - } - if item.created_at.is_none() { - item.created_at = Some(now); - } - item - }) - .collect(); - let mut turn_result = TurnResult { - turn_id: turn_id.clone(), - finish_reason: result.finish_reason.clone(), - items: output_items, - usage: result.usage.clone(), - metadata: result.metadata.clone(), - }; - if !saw_tool_call { - let editable = !matches!( - turn_result.finish_reason, - FinishReason::Cancelled | FinishReason::Error | FinishReason::Blocked - ); - self.run_turn_finish(&mut turn_result, editable, None, true) - .await?; - } - if let Some(capture) = self.telemetry.output_messages { - record_string_array_attribute( - &chat_span, - "gen_ai.output.messages", - capture_messages(&turn_result.items, capture, CaptureOrder::OldestHead), - ); - } - drop(chat_span); - self.extend_transcript(turn_result.items.clone()); - self.mark_terminal_output_committed(&turn_result.turn_id); - let output_items = turn_result.items.clone(); + } + if item.finish_reason.is_none() { + item.finish_reason = Some(finish_reason.clone()); + } + } + if item.created_at.is_none() { + item.created_at = Some(now); + } + item + }) + .collect(); + self.extend_transcript(output_items.clone()); if saw_tool_call { - self.pending_tool_batches.push(PendingToolBatch { - turn_id: turn_id.clone(), - calls: extract_tool_calls(&output_items), - results: BTreeMap::new(), - next_mutator: 0, - error: None, - }); let pending_calls = extract_tool_calls(&output_items) .into_iter() .map(|call| { @@ -3302,10 +2733,13 @@ where if let Some(step) = self.continue_active_tool_round().await? { return Ok(step); } - self.run_ready_tool_batches().await?; - self.run_turn_finish(&mut turn_result, false, None, false) - .await?; - self.finish_logical_turn(&turn_result, None).await?; + self.finish_logical_turn(&TurnResult { + turn_id, + finish_reason: result.finish_reason, + items: output_items, + usage: result.usage, + metadata: result.metadata, + }); return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput( InputRequest { session_id: self.session_id.clone(), @@ -3314,7 +2748,14 @@ where ))); } - self.finish_logical_turn(&turn_result, None).await?; + let turn_result = TurnResult { + turn_id, + finish_reason: result.finish_reason, + items: output_items, + usage: result.usage, + metadata: result.metadata, + }; + self.finish_logical_turn(&turn_result); Ok(LoopStep::Finished(turn_result)) } @@ -3403,8 +2844,7 @@ where self.append_detach_placeholder( pending.call.id.clone(), &pending.call.name, - ) - .await?; + ); } else { self.active_tool_round = Some(ActiveToolRound { presentation_turn_id: pending.presentation_turn_id.clone(), @@ -3450,53 +2890,26 @@ where } } - async fn finish_cancelled( - &mut self, - turn_id: agentkit_core::TurnId, - items: Vec, - ) -> Result { - let result = self.finish_cancelled_inner(turn_id, items).await; - if let Err(error) = &result { - self.abort_active_turn(FinishReason::Cancelled, error).await; - } - result - } - - async fn finish_cancelled_inner( + fn finish_cancelled( &mut self, turn_id: agentkit_core::TurnId, items: Vec, ) -> Result { let pending = self.drain_pending_approval_items(); self.reject_drained_approvals(pending); - let mut turn_result = TurnResult { + self.extend_transcript(items.clone()); + self.close_interrupted_tool_calls(); + let turn_result = TurnResult { turn_id, finish_reason: FinishReason::Cancelled, items, usage: None, metadata: interrupted_metadata("turn"), }; - let now = Timestamp::now(); - for item in &mut turn_result.items { - item.created_at.get_or_insert(now); - if matches!(item.kind, ItemKind::Assistant) { - item.finish_reason.get_or_insert(FinishReason::Cancelled); - } - } - self.run_turn_finish(&mut turn_result, false, None, true) - .await?; - self.extend_transcript(turn_result.items.clone()); - self.mark_terminal_output_committed(&turn_result.turn_id); - self.close_interrupted_tool_calls(); - self.run_ready_tool_batches().await?; - self.finish_logical_turn(&turn_result, None).await?; + self.finish_logical_turn(&turn_result); Ok(LoopStep::Finished(turn_result)) } - /// Buffer new input. With mutators registered, callbacks are awaited by the - /// next `next()` call before dispatch and InputAccepted is emitted only after - /// successful interception. Use `submit_input_async` for awaited acceptance. - /// /// Internal entry point for buffering user input. Reachable only via /// [`InputRequest::submit`] (resolves an `AwaitingInput` interrupt, /// including the very first one after [`Agent::start`]) and @@ -3507,72 +2920,16 @@ where /// after start-up always flows through one of the typed `submit` /// handles. pub fn submit_input(&mut self, input: Vec) -> Result<(), LoopError> { - self.check_input_submission()?; - if self.mutators.is_empty() { - self.accept_input(input); - } else { - self.pending_admission.push_back(input); - } - Ok(()) - } - - /// Await input interception before accepting this batch. On hook failure no - /// rejected input is committed and the active logical turn is unchanged. - /// Batches are processed in submission order. If an earlier queued batch is - /// rejected, this new batch remains queued for subsequent admission. - /// Sync [`Self::submit_input`] remains compatible but queues interception for - /// `next`, before dispatch, rather than awaiting acceptance in the caller. - pub async fn submit_input_async(&mut self, input: Vec) -> Result<(), LoopError> { - self.check_input_submission()?; - self.pending_admission.push_back(input); - self.admit_queued_input().await - } - - fn check_input_submission(&self) -> Result<(), LoopError> { - if self.lifecycle.closed - || self.lifecycle.closing.is_some() - || self.lifecycle.pending_cancel.is_some() - { - return Err(LoopError::InvalidState( - "session is closed or retiring a turn".into(), - )); - } if self.has_pending_interrupts() { return Err(LoopError::InvalidState( "cannot submit input while an interrupt is pending".into(), )); } - Ok(()) - } - - fn accept_input(&mut self, input: Vec) { self.emit(AgentEvent::InputAccepted { session_id: self.session_id.clone(), items: input.clone(), }); self.pending_input.extend(input); - } - - async fn intercept_input(&mut self, input: &mut Vec) -> Result<(), LoopError> { - for mutator in &self.mutators { - mutator - .on_input( - input, - self.lifecycle_ctx(self.lifecycle.active_turn.as_ref(), None), - ) - .await?; - } - Ok(()) - } - - async fn admit_queued_input(&mut self) -> Result<(), LoopError> { - while let Some(mut input) = self.pending_admission.front().cloned() { - let outcome = self.intercept_input(&mut input).await; - // Reject only this batch, preserving unrelated later submissions. - self.pending_admission.pop_front(); - outcome?; - self.accept_input(input); - } Ok(()) } @@ -3581,11 +2938,6 @@ where /// The override is consumed the next time the driver starts a model turn. /// Session-level defaults still apply to later turns. pub fn set_next_turn_cache(&mut self, cache: PromptCacheRequest) -> Result<(), LoopError> { - if self.lifecycle.closed || self.lifecycle.closing.is_some() { - return Err(LoopError::InvalidState( - "session is closed or closing".into(), - )); - } if self.has_pending_interrupts() { return Err(LoopError::InvalidState( "cannot update next-turn cache while an interrupt is pending".into(), @@ -3619,11 +2971,6 @@ where call_id: ToolCallId, decision: ApprovalDecision, ) -> Result<(), LoopError> { - if self.lifecycle.closed || self.lifecycle.closing.is_some() { - return Err(LoopError::InvalidState( - "session is closed or closing".into(), - )); - } let Some(pending) = self.pending_approvals.get_mut(&call_id) else { return Err(LoopError::InvalidState(format!( "no approval request is pending for call {}", @@ -3654,11 +3001,6 @@ where call_id: ToolCallId, input: serde_json::Value, ) -> Result<(), LoopError> { - if self.lifecycle.closed || self.lifecycle.closing.is_some() { - return Err(LoopError::InvalidState( - "session is closed or closing".into(), - )); - } let Some(pending) = self.pending_approvals.get_mut(&call_id) else { return Err(LoopError::InvalidState(format!( "no approval request is pending for call {}", @@ -3698,13 +3040,7 @@ where /// /// This clears the blocking approval and appends an error tool result so /// the transcript remains provider-valid if the host continues the turn. - /// If this retires the turn, terminal hooks run on the next `next()` or `close()`. pub fn cancel_pending_approval_for(&mut self, call_id: ToolCallId) -> Result<(), LoopError> { - if self.lifecycle.closing.is_some() || self.lifecycle.closed { - return Err(LoopError::InvalidState( - "session is closed or closing".into(), - )); - } let Some(pending) = self.drain_pending_approval_for(&call_id) else { return Err(LoopError::InvalidState(format!( "no approval request is pending for call {}", @@ -3714,24 +3050,7 @@ where let turn_id = pending.presentation_turn_id.clone(); self.reject_drained_approvals(vec![pending]); if self.pending_approvals.is_empty() && self.active_tool_round.is_none() { - if self.mutators.is_empty() { - self.close_interrupted_tool_calls(); - if self.pending_round_resume.as_ref() == Some(&turn_id) { - self.pending_round_resume = None; - } - if self.lifecycle.active_turn.as_ref() == Some(&turn_id) { - self.lifecycle.active_turn = None; - self.emit(AgentEvent::TurnFinished(TurnResult { - turn_id, - finish_reason: FinishReason::Cancelled, - items: Vec::new(), - usage: None, - metadata: interrupted_metadata("turn"), - })); - } - } else { - self.lifecycle.pending_cancel = Some(turn_id); - } + let _ = self.finish_cancelled(turn_id, Vec::new())?; } Ok(()) } @@ -3742,11 +3061,6 @@ where /// individual approval prompt. Each pending approval is resolved as denied /// and receives an error tool result so the transcript remains valid. pub async fn cancel_pending_approvals(&mut self) -> Result, LoopError> { - if self.lifecycle.closing.is_some() || self.lifecycle.closed { - return Err(LoopError::InvalidState( - "session is closed or closing".into(), - )); - } if self.pending_approvals.is_empty() { return Ok(None); } @@ -3783,10 +3097,16 @@ where self.reject_drained_approvals(pending); if let Some(error) = cleanup_error { self.close_interrupted_tool_calls(); - self.abort_active_turn(FinishReason::Error, &error).await; + self.finish_logical_turn(&TurnResult { + turn_id, + finish_reason: FinishReason::Error, + items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + }); return Err(error); } - self.finish_cancelled(turn_id, Vec::new()).await.map(Some) + self.finish_cancelled(turn_id, Vec::new()).map(Some) } /// Retire the active logical turn without resuming model or tool execution. @@ -3806,11 +3126,6 @@ where /// emitting the cancelled terminal event. The turn is retired even on error; /// a failing task manager may leave external tasks running. pub async fn retire_interrupted_turn(&mut self) -> Result, LoopError> { - if self.lifecycle.closing.is_some() || self.lifecycle.closed { - return Err(LoopError::InvalidState( - "session is closed or closing".into(), - )); - } let Some(turn_id) = self.lifecycle.active_turn.clone() else { return Ok(None); }; @@ -3818,25 +3133,15 @@ where // At AfterToolResult the foreground task round has already completed. self.pending_round_resume = None; let cleanup = self.cleanup_interrupted_turn().await; - let mut result = TurnResult { + let result = TurnResult { turn_id, finish_reason: FinishReason::Cancelled, items: Vec::new(), usage: None, metadata: interrupted_metadata("turn"), }; - self.lifecycle.pending_cancel = None; - let finish = self - .run_turn_finish(&mut result, false, cleanup.as_ref().err(), false) - .await; - let batch = self.run_ready_tool_batches().await; - let end = self - .finish_logical_turn(&result, cleanup.as_ref().err()) - .await; + self.finish_logical_turn(&result); cleanup?; - finish?; - batch?; - end?; Ok(Some(result)) } @@ -3845,12 +3150,7 @@ where LoopSnapshot { session_id: self.session_id.clone(), transcript: self.transcript.clone(), - pending_input: self - .pending_input - .iter() - .cloned() - .chain(self.pending_admission.iter().flatten().cloned()) - .collect(), + pending_input: self.pending_input.clone(), } } @@ -3896,240 +3196,45 @@ where /// Returns [`LoopError::InvalidState`] if called while an unresolved /// interrupt is pending, or propagates provider / tool / compaction errors. pub async fn next(&mut self) -> Result { - if self.lifecycle.closed || self.lifecycle.closing.is_some() { - return Err(LoopError::InvalidState( - "session is closed or closing".into(), - )); - } - let result = async { - if self.lifecycle.terminal.is_some() { - return self - .resume_terminal_delivery() - .await - .map(LoopStep::Finished); - } - if let Some(turn_id) = self.lifecycle.pending_cancel.take() { - return self.finish_cancelled(turn_id, Vec::new()).await; - } - self.admit_queued_input().await?; - if self.lifecycle.active_turn.is_none() { - let continuation_turn = self - .pending_approval_order - .iter() - .find_map(|call_id| self.pending_approvals.get(call_id)) - .map(|pending| pending.presentation_turn_id.clone()) - .or_else(|| { - self.active_tool_round - .as_ref() - .map(|active| active.presentation_turn_id.clone()) - }) - .or_else(|| self.pending_round_resume.clone()); - if let Some(turn_id) = continuation_turn { - self.start_logical_turn_with(turn_id).await?; - } else if !self.pending_input.is_empty() { - self.start_logical_turn().await?; - } - } - let step = self.next_inner().await?; - self.run_ready_tool_batches().await?; - Ok(step) - } - .await; - if let Err(error) = &result { - let reason = if matches!(error, LoopError::Cancelled) { - FinishReason::Cancelled - } else { - FinishReason::Error - }; - self.abort_active_turn(reason, error).await; - } - result - } - - async fn abort_active_turn(&mut self, reason: FinishReason, failure: &LoopError) { - self.lifecycle.pending_cancel = None; - if let Some(turn_id) = self.lifecycle.active_turn.clone() { - if let Err(error) = self.cleanup_interrupted_turn().await { - tracing::debug!(%error, "failed to clean up turn after loop error"); - } - let mut result = TurnResult { - turn_id, - finish_reason: reason, - items: Vec::new(), - usage: None, - metadata: MetadataMap::new(), - }; - if let Some(state) = &mut self.lifecycle.terminal { - state.result = result.clone(); - state.output_committed = true; - state.editable = false; - state.failure = Some(failure.clone()); - } - if !self.lifecycle.finish_attempted { - if let Err(error) = self - .run_turn_finish(&mut result, false, Some(failure), false) - .await - { - tracing::debug!(%error, "terminal finish callback failed"); - } - } - if let Err(error) = self.run_ready_tool_batches().await { - tracing::debug!(%error, "terminal batch callback failed"); - } - if let Err(error) = self.finish_logical_turn(&result, Some(failure)).await { - tracing::debug!(%error, "terminal end callback failed"); + if self.lifecycle.active_turn.is_none() { + let continuation_turn = self + .pending_approval_order + .iter() + .find_map(|call_id| self.pending_approvals.get(call_id)) + .map(|pending| pending.presentation_turn_id.clone()) + .or_else(|| { + self.active_tool_round + .as_ref() + .map(|active| active.presentation_turn_id.clone()) + }) + .or_else(|| self.pending_round_resume.clone()); + if let Some(turn_id) = continuation_turn { + self.start_logical_turn_with(turn_id); + } else if !self.pending_input.is_empty() { + self.start_logical_turn(); } } - } - /// Explicitly close this session, retiring an active turn as cancelled. - /// If final delivery already began in `next`, its existing candidate, - /// outcome and callback progress are resumed instead of being replaced. - /// Admission and `next` are disabled as soon as closing starts. Dropping this - /// future leaves resumable progress: a later `close` retries the in-flight - /// callback, but never repeats completed callbacks or the terminal event. - /// Callbacks must tolerate cancellation/retry of their own unfinished work. - /// All cleanup/end callbacks are attempted, with the first error returned. - /// After completion (even on error), repeated closes are no-ops. Detached - /// background work is not awaited; Drop does not run callbacks. - pub async fn close(&mut self) -> Result<(), LoopError> { - if self.lifecycle.closed { - return Ok(()); - } - if self.lifecycle.closing.is_none() { - let mut cleanup_turns = Vec::new(); - if let Some(active) = &self.active_tool_round { - cleanup_turns.push(active.task_turn_id.clone()); - } - if let Some(turn_id) = &self.pending_round_resume { - if !cleanup_turns.contains(turn_id) { - cleanup_turns.push(turn_id.clone()); - } - } - for pending in self.pending_approvals.values() { - if !cleanup_turns.contains(&pending.tool_request.turn_id) { - cleanup_turns.push(pending.tool_request.turn_id.clone()); - } - } - if self.lifecycle.terminal.is_none() { + let result = self.next_inner().await; + match &result { + Ok(LoopStep::Finished(turn)) => self.finish_logical_turn(turn), + Err(_) => { if let Some(turn_id) = self.lifecycle.active_turn.clone() { - let result = TurnResult { + if let Err(error) = self.cleanup_interrupted_turn().await { + tracing::debug!(%error, "failed to clean up turn after loop error"); + } + self.finish_logical_turn(&TurnResult { turn_id, - finish_reason: FinishReason::Cancelled, + finish_reason: FinishReason::Error, items: Vec::new(), usage: None, - metadata: interrupted_metadata("turn"), - }; - self.begin_terminal(&result, false, None, false); + metadata: MetadataMap::new(), + }); } } - self.lifecycle.closing = Some(ClosingState { - cleanup_turns, - cleanup_index: 0, - cleanup_done: false, - session_index: 0, - error: None, - }); - } - loop { - let state = self.lifecycle.closing.as_ref().expect("closing state"); - if state.cleanup_index == state.cleanup_turns.len() { - break; - } - let turn_id = state.cleanup_turns[state.cleanup_index].clone(); - let outcome = self.task_manager.on_turn_interrupted(&turn_id).await; - let state = self.lifecycle.closing.as_mut().expect("closing state"); - state.cleanup_index += 1; - if let Err(error) = outcome { - state - .error - .get_or_insert(LoopError::Tool(ToolError::Internal(error.to_string()))); - } - } - if !self - .lifecycle - .closing - .as_ref() - .expect("closing state") - .cleanup_done - { - self.active_tool_round = None; - self.pending_round_resume = None; - let pending = self.drain_pending_approval_items(); - self.reject_drained_approvals(pending); - self.close_interrupted_tool_calls(); - self.pending_input.clear(); - self.pending_admission.clear(); - self.lifecycle.pending_cancel = None; - self.lifecycle - .closing - .as_mut() - .expect("closing state") - .cleanup_done = true; - } - if self.lifecycle.terminal.is_some() { - let cleanup_error = self - .lifecycle - .closing - .as_ref() - .expect("closing state") - .error - .clone(); - if let Some(error) = cleanup_error { - self.lifecycle - .terminal - .as_mut() - .expect("terminal state") - .failure - .get_or_insert(error); - } - if let Err(error) = self.resume_terminal_delivery().await { - self.lifecycle - .closing - .as_mut() - .expect("closing state") - .error - .get_or_insert(error); - } - } else if let Err(error) = self.run_ready_tool_batches().await { - self.lifecycle - .closing - .as_mut() - .expect("closing state") - .error - .get_or_insert(error); - } - loop { - let state = self.lifecycle.closing.as_ref().expect("closing state"); - if state.session_index == self.mutators.len() { - break; - } - let outcome = self.mutators[state.session_index] - .on_session_end( - SessionEndPayload { - transcript: &self.transcript, - }, - LifecycleCtx { - error: state.error.as_ref(), - ..self.lifecycle_ctx(None, None) - }, - ) - .await; - let state = self.lifecycle.closing.as_mut().expect("closing state"); - state.session_index += 1; - if let Err(error) = outcome { - state.error.get_or_insert(error); - } + _ => {} } - self.session = None; - self.pending_tool_batches.clear(); - self.lifecycle.closed = true; - self.lifecycle - .closing - .take() - .expect("closing state") - .error - .map_or(Ok(()), Err) + result } async fn cleanup_interrupted_turn(&mut self) -> Result<(), LoopError> { @@ -4194,14 +3299,11 @@ where // before unrelated background completions so a delayed approval cannot // bind itself to that turn's TurnStarted event. AfterToolResult resumes // remain ordered ahead of fresh input below. - if self.pending_round_resume.is_none() - && (!self.pending_input.is_empty() || self.lifecycle.started_with_input) - { - self.lifecycle.started_with_input = false; + if self.pending_round_resume.is_none() && !self.pending_input.is_empty() { // Take updates now to preserve the driver's once-per-step manager // handoff, but defer presenting them until this input turn ends. self.collect_pending_loop_updates().await?; - let turn_id = self.start_logical_turn().await?; + let turn_id = self.start_logical_turn(); let drained: Vec = std::mem::take(&mut self.pending_input); self.extend_transcript(drained); return self @@ -4235,7 +3337,7 @@ where ))); } - let turn_id = self.start_logical_turn().await?; + let turn_id = self.start_logical_turn(); let drained: Vec = std::mem::take(&mut self.pending_input); self.extend_transcript(drained); self.drive_turn(turn_id, MutationPoint::AfterTurnEnded) @@ -4250,31 +3352,7 @@ where /// registered [`TranscriptObserver`]s. The single mutation point — /// every push to `self.transcript` should funnel through here so /// observers see exactly what landed in the transcript. - fn commit_transcript(&mut self, mut candidate: Vec) { - if candidate.starts_with(&self.transcript) { - let suffix = candidate.split_off(self.transcript.len()); - self.extend_transcript(suffix); - } else { - let now = Timestamp::now(); - for item in &mut candidate { - item.created_at.get_or_insert(now); - } - self.transcript = candidate; - for observer in &self.transcript_observers { - observer.on_transcript_rewrite(TranscriptRewriteEvent { - session_id: &self.session_id, - items: &self.transcript, - }); - } - } - } - fn append_item(&mut self, mut item: Item) { - for part in &item.parts { - if let Part::ToolResult(result) = part { - self.record_batch_result(result); - } - } if item.created_at.is_none() { item.created_at = Some(Timestamp::now()); } @@ -4287,13 +3365,10 @@ where self.transcript.push(item); } - async fn append_detach_placeholder( - &mut self, - call_id: ToolCallId, - tool_name: &str, - ) -> Result<(), LoopError> { - if self.detached_call_ids.contains(&call_id) { - return Ok(()); + fn append_detach_placeholder(&mut self, call_id: ToolCallId, tool_name: &str) { + self.background_call_ids.insert(call_id.clone()); + if !self.detached_call_ids.insert(call_id.clone()) { + return; } let detached_result = ToolResultPart { call_id: call_id.clone(), @@ -4303,24 +3378,6 @@ where is_error: false, metadata: MetadataMap::new(), }; - for mutator in &self.mutators { - mutator - .on_tool_progress( - &detached_result, - self.lifecycle_ctx(self.lifecycle.active_turn.as_ref(), None), - ) - .await?; - } - for mutator in &self.mutators { - mutator - .on_turn_progress( - TurnProgress::ToolDetached(&detached_result), - self.lifecycle_ctx(self.lifecycle.active_turn.as_ref(), None), - ) - .await?; - } - self.background_call_ids.insert(call_id.clone()); - self.detached_call_ids.insert(call_id); self.emit(AgentEvent::ToolExecutionProgress(detached_result.clone())); self.append_item(Item { id: None, @@ -4331,7 +3388,6 @@ where finish_reason: None, created_at: None, }); - Ok(()) } /// Append a tool-result Item: emit one [`AgentEvent::ToolResultReceived`] @@ -5136,7 +4192,7 @@ use std::collections::VecDeque; use std::sync::{Arc, Mutex}; use std::time::SystemTime; use agentkit_core::{CostUsage, DataRef, FinishReason, Item, ItemKind, MetadataMap, Modality, Part, TokenUsage, TurnCancellation, Usage}; -use agentkit_loop::{Agent, LoopError, LoopStep, MessageCapture, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, TelemetryConfig, TurnRequest, set_provider_finish_reasons}; +use agentkit_loop::{Agent, HookCtx, LoopError, LoopMutator, LoopStep, MessageCapture, ModelRequestPayload, ModelResponsePayload, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, TelemetryConfig, TurnRequest, set_provider_finish_reasons}; use async_trait::async_trait; use opentelemetry::trace::{Span, SpanBuilder, SpanContext, Status, Tracer}; use opentelemetry::{Array, Context, KeyValue, Value}; @@ -5240,7 +4296,23 @@ fn json_array(span: &Exported, key: &str) -> Vec { } } +struct EditMessages; +#[async_trait] +impl LoopMutator for EditMessages { + async fn on_model_request(&self, payload: ModelRequestPayload<'_>, _: HookCtx<'_>) -> Result<(), LoopError> { + if let Some(Part::Text(text)) = payload.transcript.last_mut().and_then(|item| item.parts.first_mut()) { text.text.push_str("-edited"); } + Ok(()) + } + async fn on_model_response(&self, payload: ModelResponsePayload<'_>, _: HookCtx<'_>) -> Result<(), LoopError> { + for item in payload.output { for part in item.parts.iter_mut() { if let Part::Text(text) = part { text.text.push_str("-edited"); } } } + Ok(()) + } +} + fn run_attempt(adapter: ScriptedAdapter) -> (Vec, Result) { + run_attempt_with(adapter, false) +} +fn run_attempt_with(adapter: ScriptedAdapter, edit: bool) -> (Vec, Result) { let tracer = MemoryTracer::default(); let subscriber = tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer.clone())); let runtime = tokio::runtime::Builder::new_current_thread().build().unwrap(); @@ -5249,13 +4321,14 @@ fn run_attempt(adapter: ScriptedAdapter) -> (Vec, Result result, other => panic!("unexpected step: {other:?}") }; + let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap(); + let input = json_array(chat, "gen_ai.input.messages"); + assert!(input.last().unwrap().to_string().contains("newest-edited")); + let output = json_array(chat, "gen_ai.output.messages"); + assert!(output[0].to_string().contains("first-edited")); + assert!(output[1].to_string().contains("second-edited")); + assert!(matches!(&result.items[0].parts[0], Part::Text(text) if text.text == "first-edited")); } "#; } @@ -5917,7 +5000,7 @@ fn tool_result_not_started(item: &Item) -> bool { } /// Errors that can occur while driving the agent loop. -#[derive(Debug, Clone, Error)] +#[derive(Debug, Error)] pub enum LoopError { /// Typed, sanitized model failure with retry accounting. #[error(transparent)] @@ -5990,102 +5073,6 @@ fn fan_out_observed_event( /// The only invariant currently checked is tool_use ↔ tool_result pairing /// — every [`Part::ToolCall`] must be followed (in transcript order) by a /// matching [`Part::ToolResult`] with the same `call_id`. -fn validate_prompt_edits(original: &[Item], candidate: &[Item]) -> Result<(), LoopError> { - validate_transcript_invariants(candidate)?; - let identified = original.iter().filter(|item| item.id.is_some()); - for original_item in identified { - let matches: Vec<&Item> = candidate - .iter() - .filter(|item| item.id == original_item.id) - .collect(); - if matches.len() != 1 - || matches[0].kind != original_item.kind - || matches[0].usage != original_item.usage - || matches[0].finish_reason != original_item.finish_reason - || matches[0].created_at != original_item.created_at - { - return Err(LoopError::Mutator( - "prompt hook changed a message identity or accounting".into(), - )); - } - } - for original_item in original.iter().filter(|item| { - item.id.is_none() - && (item.usage.is_some() || item.finish_reason.is_some() || item.created_at.is_some()) - }) { - let same_facts = |item: &&Item| { - item.id.is_none() - && item.kind == original_item.kind - && item.usage == original_item.usage - && item.finish_reason == original_item.finish_reason - && item.created_at == original_item.created_at - }; - if candidate.iter().filter(same_facts).count() < original.iter().filter(same_facts).count() - { - return Err(LoopError::Mutator( - "prompt hook removed or changed anonymous item accounting".into(), - )); - } - } - let tool_parts = |items: &[Item]| -> Vec { - items - .iter() - .flat_map(|item| &item.parts) - .filter(|part| matches!(part, Part::ToolCall(_) | Part::ToolResult(_))) - .cloned() - .collect() - }; - if tool_parts(original) != tool_parts(candidate) { - return Err(LoopError::Mutator( - "prompt hook changed existing tool calls or results".into(), - )); - } - Ok(()) -} - -fn output_payload<'a>( - original: &'a [Item], - candidate: &'a mut [Item], -) -> Vec> { - original - .iter() - .zip(candidate) - .map(|(original, item)| OutputItemPayload { - original, - parts: &mut item.parts, - metadata: &mut item.metadata, - }) - .collect() -} - -fn validate_output_edits(original: &[Item], candidate: &[Item]) -> Result<(), LoopError> { - fn links(item: &Item) -> Vec<(bool, ToolCallId, Option)> { - item.parts - .iter() - .filter_map(|part| match part { - Part::ToolCall(call) => Some((true, call.id.clone(), Some(call.name.clone()))), - Part::ToolResult(result) => Some((false, result.call_id.clone(), None)), - _ => None, - }) - .collect() - } - if original.len() != candidate.len() - || original.iter().zip(candidate).any(|(a, b)| { - a.id != b.id - || a.kind != b.kind - || a.usage != b.usage - || a.finish_reason != b.finish_reason - || a.created_at != b.created_at - || links(a) != links(b) - }) - { - return Err(LoopError::Mutator( - "output hook changed immutable item identity, accounting, or tool linkage".into(), - )); - } - Ok(()) -} - fn validate_transcript_invariants(transcript: &[Item]) -> Result<(), LoopError> { let mut pending: HashSet = HashSet::new(); let mut seen_calls: HashSet = HashSet::new(); @@ -6130,6 +5117,98 @@ fn validate_transcript_invariants(transcript: &[Item]) -> Result<(), LoopError> Ok(()) } +async fn run_request_hooks( + mutators: &[Arc], + session_id: &SessionId, + turn_id: &agentkit_core::TurnId, + cancellation: Option<&TurnCancellation>, + request: &mut TurnRequest, +) -> Result<(), LoopError> { + if mutators.is_empty() { + return Ok(()); + } + for mutator in mutators { + let payload = ModelRequestPayload { + transcript: &mut request.transcript, + available_tools: &mut request.available_tools, + cache: &mut request.cache, + metadata: &mut request.metadata, + }; + let ctx = HookCtx { + session_id, + turn_id: Some(turn_id), + cancellation: cancellation.cloned(), + }; + mutator.on_model_request(payload, ctx).await?; + } + validate_transcript_invariants(&request.transcript) +} + +async fn run_response_hooks( + mutators: &[Arc], + session_id: &SessionId, + turn_id: &agentkit_core::TurnId, + cancellation: Option<&TurnCancellation>, + result: &mut ModelTurnResult, +) -> Result<(), LoopError> { + if mutators.is_empty() { + return Ok(()); + } + let original = result.output_items.clone(); + for mutator in mutators { + let entry = result.output_items.clone(); + let payload = ModelResponsePayload { + output: entry + .iter() + .zip(&mut result.output_items) + .map(|(original, item)| OutputItemPayload { + original, + parts: &mut item.parts, + metadata: &mut item.metadata, + }) + .collect(), + finish_reason: &result.finish_reason, + usage: result.usage.as_ref(), + metadata: &result.metadata, + }; + let ctx = HookCtx { + session_id, + turn_id: Some(turn_id), + cancellation: cancellation.cloned(), + }; + mutator.on_model_response(payload, ctx).await?; + } + validate_output_edits(&original, &result.output_items) +} + +/// Hard-fails when a response hook changed anything other than item content: +/// item count, identity, role, accounting, or tool-call linkage. +fn validate_output_edits(original: &[Item], edited: &[Item]) -> Result<(), LoopError> { + fn tool_links(item: &Item) -> impl Iterator)> { + item.parts.iter().filter_map(|part| match part { + Part::ToolCall(call) => Some((&call.id, Some(call.name.as_str()))), + Part::ToolResult(result) => Some((&result.call_id, None)), + _ => None, + }) + } + let unchanged = original.len() == edited.len() + && original.iter().zip(edited).all(|(a, b)| { + a.id == b.id + && a.kind == b.kind + && a.usage == b.usage + && a.finish_reason == b.finish_reason + && a.created_at == b.created_at + && tool_links(a).eq(tool_links(b)) + }); + if unchanged { + Ok(()) + } else { + Err(LoopError::Mutator( + "response hook changed item identity, accounting, or tool-call linkage".into(), + )) + } +} + #[cfg(test)] mod tests { use std::collections::VecDeque; @@ -9991,7 +9070,6 @@ mod tests { )); driver .finish_cancelled(agentkit_core::TurnId::new("turn-cancelled"), Vec::new()) - .await .unwrap(); assert!( @@ -11299,4 +10377,4 @@ mod tests { } #[cfg(test)] -mod lifecycle_tests; +mod hook_tests; diff --git a/crates/agentkit-loop/src/lifecycle_tests.rs b/crates/agentkit-loop/src/lifecycle_tests.rs deleted file mode 100644 index 24812bc..0000000 --- a/crates/agentkit-loop/src/lifecycle_tests.rs +++ /dev/null @@ -1,1415 +0,0 @@ -use super::*; -use agentkit_tools_core::{ToolContext, ToolExecutionOutcome, ToolResult}; -use std::sync::{ - Mutex, - atomic::{AtomicBool, Ordering}, -}; - -#[derive(Clone, Copy)] -enum Script { - Reply, - Tool, - BeginError, - StreamError, - MissingFinished, - Cancelled, -} - -#[derive(Clone)] -struct ProbeModel { - scripts: Arc>>, - requests: Arc>>, - configs: Arc>>, -} -impl ProbeModel { - fn new(scripts: impl IntoIterator) -> Self { - Self { - scripts: Arc::new(Mutex::new(scripts.into_iter().collect())), - requests: Arc::default(), - configs: Arc::default(), - } - } -} -struct ProbeSession(ProbeModel); -struct ProbeTurn(VecDeque>); -#[async_trait] -impl ModelAdapter for ProbeModel { - type Session = ProbeSession; - async fn start_session(&self, config: SessionConfig) -> Result { - self.configs.lock().unwrap().push(config); - Ok(ProbeSession(self.clone())) - } -} -#[async_trait] -impl ModelSession for ProbeSession { - type Turn = ProbeTurn; - async fn begin_turn( - &mut self, - request: TurnRequest, - _: Option, - ) -> Result { - self.0.requests.lock().unwrap().push(request); - let script = self - .0 - .scripts - .lock() - .unwrap() - .pop_front() - .expect("unexpected inference"); - let item = match script { - Script::BeginError => return Err(LoopError::Provider("original begin error".into())), - Script::Cancelled => return Err(LoopError::Cancelled), - Script::StreamError => { - return Ok(ProbeTurn(VecDeque::from([Err(LoopError::Provider( - "original stream error".into(), - ))]))); - } - Script::MissingFinished => return Ok(ProbeTurn(VecDeque::new())), - Script::Reply => Item::text(ItemKind::Assistant, "reply").with_id("message-reply"), - Script::Tool => Item::new( - ItemKind::Assistant, - vec![Part::ToolCall(ToolCallPart::new( - "call-probe", - "probe", - serde_json::json!({"value": 1}), - ))], - ) - .with_id("message-tool"), - }; - let finish_reason = if matches!(script, Script::Tool) { - FinishReason::ToolCall - } else { - FinishReason::Completed - }; - let mut events = VecDeque::new(); - if let Some(Part::ToolCall(call)) = item.parts.first() { - events.push_back(Ok(ModelTurnEvent::ToolCall(call.clone()))); - } - events.push_back(Ok(ModelTurnEvent::Finished(ModelTurnResult { - finish_reason, - output_items: vec![item], - usage: None, - metadata: MetadataMap::new(), - model: Some("probe-model".into()), - response_id: Some("response-probe".into()), - }))); - Ok(ProbeTurn(events)) - } -} -#[async_trait] -impl ModelTurn for ProbeTurn { - async fn next_event( - &mut self, - _: Option, - ) -> Result, LoopError> { - self.0.pop_front().transpose() - } -} - -struct ProbeExecutor; -#[async_trait] -impl ToolExecutor for ProbeExecutor { - fn specs(&self) -> Vec { - Vec::new() - } - async fn execute(&self, request: ToolRequest, _: &mut ToolContext<'_>) -> ToolExecutionOutcome { - ToolExecutionOutcome::Completed(ToolResult { - result: ToolResultPart { - call_id: request.call_id, - output: ToolOutput::Text("tool-result".into()), - is_error: false, - metadata: MetadataMap::new(), - }, - duration: None, - metadata: MetadataMap::new(), - }) - } -} - -#[derive(Clone)] -struct ProbeHook { - label: &'static str, - events: Arc>>, - fail: Option<&'static str>, - failed: Arc, - break_links: bool, -} -impl ProbeHook { - fn new(label: &'static str, events: &Arc>>) -> Self { - Self { - label, - events: events.clone(), - fail: None, - failed: Arc::new(AtomicBool::new(false)), - break_links: false, - } - } - fn event(&self, name: &str, ctx: &LifecycleCtx<'_>) -> Result<(), LoopError> { - self.events.lock().unwrap().push(format!( - "{}:{name}:{}:{}", - self.label, - ctx.turn_id.map(|id| id.0.as_str()).unwrap_or("-"), - ctx.model_call_index - .map(|i| i.to_string()) - .unwrap_or_else(|| "-".into()) - )); - if self.fail == Some(name) && !self.failed.swap(true, Ordering::SeqCst) { - Err(LoopError::Mutator(format!("{} failed {name}", self.label))) - } else { - Ok(()) - } - } - fn rewrite(&self, output: &mut [OutputItemPayload<'_>], suffix: &str) { - for item in output { - for part in item.parts.iter_mut() { - if let Part::Text(text) = part { - text.text.push_str(&format!(":{suffix}{}", self.label)); - } - } - } - } -} -#[async_trait] -impl LoopMutator for ProbeHook { - async fn on_session_start( - &self, - payload: SessionStartPayload<'_>, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - payload.metadata.insert(self.label.into(), true.into()); - self.event("session_start", &ctx) - } - async fn on_session_end( - &self, - _: SessionEndPayload<'_>, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - self.event("session_end", &ctx) - } - async fn on_input( - &self, - input: &mut Vec, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - for item in input { - for part in &mut item.parts { - if let Part::Text(text) = part { - text.text.push_str(self.label); - } - } - } - self.event("input", &ctx) - } - async fn on_turn_start( - &self, - _: &mut TranscriptCursor<'_>, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - self.event("turn_start", &ctx) - } - async fn on_turn_finish( - &self, - mut payload: TurnFinishPayload<'_>, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - if let Some(output) = &mut payload.output { - self.rewrite(output, "f"); - } - if matches!( - payload.result.finish_reason, - FinishReason::Error | FinishReason::Cancelled - ) { - assert!(payload.output.is_none()); - } - self.event("turn_finish", &ctx) - } - async fn on_turn_end(&self, _: &TurnResult, ctx: LifecycleCtx<'_>) -> Result<(), LoopError> { - self.event("turn_end", &ctx) - } - async fn on_model_request( - &self, - payload: ModelRequestPayload<'_>, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - payload.metadata.insert(self.label.into(), true.into()); - self.event("request", &ctx) - } - async fn on_model_response( - &self, - mut payload: ModelResponsePayload<'_>, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - self.rewrite(&mut payload.output, "m"); - if self.break_links { - for item in payload.output { - for part in item.parts.iter_mut() { - if let Part::ToolCall(call) = part { - call.id = ToolCallId::new("forged"); - } - } - } - } - self.event("response", &ctx) - } - async fn on_model_progress( - &self, - _: &ModelTurnEvent, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - self.event("progress", &ctx) - } - async fn on_model_error(&self, _: &LoopError, ctx: LifecycleCtx<'_>) -> Result<(), LoopError> { - self.event("model_error", &ctx) - } - async fn on_tool_batch( - &self, - payload: ToolBatchPayload<'_>, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - assert_eq!(payload.calls.len(), 1); - assert_eq!(payload.results.len(), 1); - assert_eq!(payload.calls[0].id, payload.results[0].call_id); - self.event("batch", &ctx) - } -} -struct Capture(Arc>>); -impl LoopObserver for Capture { - fn handle_event(&self, event: ObservedEvent) { - self.0.lock().unwrap().push(event.event); - } -} -fn count(events: &Arc>>, name: &str) -> usize { - events - .lock() - .unwrap() - .iter() - .filter(|event| event.contains(&format!(":{name}:"))) - .count() -} -fn text(item: &Item) -> &str { - match &item.parts[0] { - Part::Text(text) => &text.text, - _ => panic!("expected text"), - } -} -fn finished(step: LoopStep) -> TurnResult { - match step { - LoopStep::Finished(result) => result, - other => panic!("expected finished, got {other:?}"), - } -} - -#[tokio::test] -async fn lifecycle_sequential_payloads_and_committed_final_output() { - let events = Arc::default(); - let observed: Arc>> = Arc::default(); - let model = ProbeModel::new([Script::Reply]); - let agent = Agent::builder() - .model(model.clone()) - .input(vec![Item::text(ItemKind::User, "input")]) - .mutator(ProbeHook::new("1", &events)) - .mutator(ProbeHook::new("2", &events)) - .observer(Capture(observed.clone())) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("lifecycle")).await.unwrap(); - assert_eq!(model.configs.lock().unwrap()[0].metadata.len(), 2); - let result = finished(driver.next().await.unwrap()); - let request = &model.requests.lock().unwrap()[0]; - assert_eq!(request.session_id, SessionId::new("lifecycle")); - assert_eq!(text(&request.transcript[0]), "input12"); - assert_eq!(request.metadata.len(), 2); - assert_eq!(text(&result.items[0]), "reply:m1:m2:f1:f2"); - assert_eq!(driver.snapshot().transcript.last(), result.items.last()); - assert!( - observed - .lock() - .unwrap() - .iter() - .any(|event| matches!(event, AgentEvent::TurnFinished(turn) if turn == &result)) - ); - let calls = events.lock().unwrap(); - for pair in calls.chunks_exact(2) { - assert!(pair[0].starts_with("1:")); - assert_eq!(&pair[0][1..], &pair[1][1..]); - } - drop(calls); - assert_eq!(count(&events, "turn_start"), 2); - assert_eq!(count(&events, "turn_finish"), 2); - assert_eq!(count(&events, "turn_end"), 2); - assert_eq!(count(&events, "session_end"), 0); - driver.close().await.unwrap(); - driver.close().await.unwrap(); - assert_eq!(count(&events, "session_end"), 2); - assert!(matches!( - driver.next().await, - Err(LoopError::InvalidState(_)) - )); - assert!(matches!( - driver.submit_input(vec![]), - Err(LoopError::InvalidState(_)) - )); - assert!(matches!( - driver.submit_input_async(vec![]).await, - Err(LoopError::InvalidState(_)) - )); -} - -#[tokio::test] -async fn sync_admission_is_pre_dispatch_and_async_admission_is_pre_acceptance() { - let events = Arc::default(); - let observed: Arc>> = Arc::default(); - let agent = Agent::builder() - .model(ProbeModel::new([Script::Reply, Script::Reply])) - .mutator(ProbeHook::new("1", &events)) - .observer(Capture(observed.clone())) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("input")).await.unwrap(); - driver - .submit_input(vec![Item::text(ItemKind::User, "sync")]) - .unwrap(); - assert_eq!(count(&events, "input"), 0); - assert_eq!(driver.snapshot().pending_input.len(), 1); - assert!( - !observed - .lock() - .unwrap() - .iter() - .any(|event| matches!(event, AgentEvent::InputAccepted { .. })) - ); - finished(driver.next().await.unwrap()); - assert_eq!(count(&events, "input"), 1); - driver - .submit_input_async(vec![Item::text(ItemKind::User, "async")]) - .await - .unwrap(); - assert_eq!(count(&events, "input"), 2); - finished(driver.next().await.unwrap()); - assert_eq!(count(&events, "input"), 2); -} - -#[tokio::test] -async fn input_failure_drops_admission_without_starting_a_turn_and_can_recover() { - for awaited in [false, true] { - let events = Arc::default(); - let model = ProbeModel::new([Script::Reply]); - let mut hook = ProbeHook::new("1", &events); - hook.fail = Some("input"); - let agent = Agent::builder() - .model(model.clone()) - .mutator(hook) - .build() - .unwrap(); - let mut driver = agent - .start(SessionConfig::new("input-failure")) - .await - .unwrap(); - let input = vec![Item::text(ItemKind::User, "rejected")]; - if awaited { - assert!(driver.submit_input_async(input).await.is_err()); - } else { - driver.submit_input(input).unwrap(); - assert!(driver.next().await.is_err()); - } - assert!(driver.snapshot().transcript.is_empty()); - assert!(driver.snapshot().pending_input.is_empty()); - assert_eq!(count(&events, "turn_start"), 0); - assert!(model.requests.lock().unwrap().is_empty()); - driver - .submit_input_async(vec![Item::text(ItemKind::User, "accepted")]) - .await - .unwrap(); - finished(driver.next().await.unwrap()); - } -} - -#[tokio::test] -async fn tool_continuation_keeps_one_logical_turn_and_reports_real_batch() { - let events = Arc::default(); - let model = ProbeModel::new([Script::Tool, Script::Reply]); - let agent = Agent::builder() - .model(model.clone()) - .tool_executor(ProbeExecutor) - .mutator(ProbeHook::new("1", &events)) - .input(vec![Item::text(ItemKind::User, "tools")]) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("tools")).await.unwrap(); - let turn_id = match driver.next().await.unwrap() { - LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info.turn_id, - other => panic!("{other:?}"), - }; - assert_eq!(count(&events, "batch"), 1); - assert_eq!(count(&events, "turn_end"), 0); - driver - .submit_input(vec![Item::text(ItemKind::User, "interject")]) - .unwrap(); - let result = finished(driver.next().await.unwrap()); - assert_eq!(result.turn_id, turn_id); - assert_eq!(count(&events, "turn_start"), 1); - assert_eq!(count(&events, "turn_finish"), 1); - assert_eq!(count(&events, "turn_end"), 1); - let requests = model.requests.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert_eq!(requests[0].turn_id, requests[1].turn_id); - let calls = events.lock().unwrap(); - assert!( - calls - .iter() - .any(|event| event.ends_with(":1") && event.contains(":request:")) - ); - assert!( - calls - .iter() - .any(|event| event.ends_with(":2") && event.contains(":request:")) - ); -} - -#[tokio::test] -async fn output_hook_failure_is_atomic_and_finishes_once() { - for point in [ - "response", - "turn_finish", - "progress", - "request", - "turn_start", - ] { - let events = Arc::default(); - let observed: Arc>> = Arc::default(); - let mut hook = ProbeHook::new("1", &events); - hook.fail = Some(point); - let agent = Agent::builder() - .model(ProbeModel::new([Script::Reply])) - .mutator(hook) - .input(vec![Item::text(ItemKind::User, "input")]) - .observer(Capture(observed.clone())) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("atomic")).await.unwrap(); - assert!(matches!(driver.next().await, Err(LoopError::Mutator(_)))); - assert!( - !driver - .snapshot() - .transcript - .iter() - .any(|item| item.kind == ItemKind::Assistant) - ); - assert_eq!(count(&events, "turn_finish"), 1, "{point}"); - assert_eq!(count(&events, "turn_end"), 1, "{point}"); - assert!(matches!( - driver.next().await.unwrap(), - LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) - )); - assert_eq!(count(&events, "turn_end"), 1); - assert_eq!( - observed - .lock() - .unwrap() - .iter() - .filter(|event| matches!(event, AgentEvent::TurnFinished(_))) - .count(), - 1 - ); - } -} - -#[tokio::test] -async fn model_failures_preserve_original_error_despite_terminal_hook_failures() { - for script in [ - Script::BeginError, - Script::StreamError, - Script::MissingFinished, - ] { - let events = Arc::default(); - let mut hook = ProbeHook::new("1", &events); - hook.fail = Some("model_error"); - let mut terminal = ProbeHook::new("2", &events); - terminal.fail = Some("turn_finish"); - let mut end = ProbeHook::new("3", &events); - end.fail = Some("turn_end"); - let agent = Agent::builder() - .model(ProbeModel::new([script])) - .mutator(hook) - .mutator(terminal) - .mutator(end) - .input(vec![Item::text(ItemKind::User, "input")]) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("failure")).await.unwrap(); - assert!(matches!(driver.next().await, Err(LoopError::Provider(_)))); - assert_eq!(count(&events, "model_error"), 3); - assert_eq!(count(&events, "turn_finish"), 3); - assert_eq!(count(&events, "turn_end"), 3); - driver.close().await.unwrap(); - assert_eq!(count(&events, "turn_end"), 3); - } -} - -#[tokio::test] -async fn cancellation_and_explicit_retirement_run_readonly_finish_and_end() { - let events = Arc::default(); - let agent = Agent::builder() - .model(ProbeModel::new([Script::Cancelled])) - .mutator(ProbeHook::new("1", &events)) - .input(vec![Item::text(ItemKind::User, "input")]) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("cancelled")).await.unwrap(); - let result = finished(driver.next().await.unwrap()); - assert_eq!(result.finish_reason, FinishReason::Cancelled); - assert_eq!(driver.snapshot().transcript.last(), result.items.last()); - assert_eq!(count(&events, "turn_finish"), 1); - assert_eq!(count(&events, "turn_end"), 1); - assert_eq!(count(&events, "model_error"), 1); - assert!(driver.retire_interrupted_turn().await.unwrap().is_none()); - driver.close().await.unwrap(); - assert_eq!(count(&events, "turn_end"), 1); - - let events = Arc::default(); - let agent = Agent::builder() - .model(ProbeModel::new([Script::Tool])) - .tool_executor(ProbeExecutor) - .mutator(ProbeHook::new("1", &events)) - .input(vec![Item::text(ItemKind::User, "input")]) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("retire")).await.unwrap(); - driver.next().await.unwrap(); - driver.close().await.unwrap(); - assert_eq!(count(&events, "turn_finish"), 1); - assert_eq!(count(&events, "turn_end"), 1); - assert_eq!(count(&events, "session_end"), 1); -} - -#[tokio::test] -async fn forged_tool_linkage_is_rejected_before_commit_or_execution() { - let events = Arc::default(); - let mut hook = ProbeHook::new("1", &events); - hook.break_links = true; - let agent = Agent::builder() - .model(ProbeModel::new([Script::Tool])) - .tool_executor(ProbeExecutor) - .mutator(hook) - .input(vec![Item::text(ItemKind::User, "input")]) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("identity")).await.unwrap(); - assert!(matches!(driver.next().await, Err(LoopError::Mutator(_)))); - assert_eq!(driver.snapshot().transcript.len(), 1); - assert_eq!(count(&events, "batch"), 0); - assert_eq!(count(&events, "turn_end"), 1); -} - -#[tokio::test] -async fn session_end_failure_does_not_skip_later_hooks_or_reopen_session() { - let events = Arc::default(); - let mut hook = ProbeHook::new("1", &events); - hook.fail = Some("session_end"); - let agent = Agent::builder() - .model(ProbeModel::new([])) - .mutator(hook) - .mutator(ProbeHook::new("2", &events)) - .build() - .unwrap(); - let mut driver = agent - .start(SessionConfig::new("close-error")) - .await - .unwrap(); - assert!(driver.close().await.is_err()); - assert_eq!(count(&events, "session_end"), 2); - driver.close().await.unwrap(); - assert_eq!(count(&events, "session_end"), 2); - assert!(driver.next().await.is_err()); -} - -struct PromptHook { - fail_start: bool, - fail_request: bool, -} -#[async_trait] -impl LoopMutator for PromptHook { - async fn on_turn_start( - &self, - cursor: &mut TranscriptCursor<'_>, - _: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - for item in cursor.iter_mut() { - if item.kind == ItemKind::User { - item.parts = vec![Part::text("logical prompt")]; - } - } - cursor.insert(0, Item::text(ItemKind::System, "persistent context")); - if self.fail_start { - Err(LoopError::Mutator("start prompt failed".into())) - } else { - Ok(()) - } - } - async fn on_model_request( - &self, - payload: ModelRequestPayload<'_>, - _: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - for item in payload.transcript.iter_mut() { - if item.kind == ItemKind::User { - item.parts = vec![Part::text("inference prompt")]; - } - } - payload - .transcript - .insert(0, Item::text(ItemKind::System, "inference-only context")); - if self.fail_request { - Err(LoopError::Mutator("request prompt failed".into())) - } else { - Ok(()) - } - } -} - -#[tokio::test] -async fn logical_prompt_edits_persist_but_inference_prompt_edits_do_not() { - let model = ProbeModel::new([Script::Reply]); - let agent = Agent::builder() - .model(model.clone()) - .mutator(PromptHook { - fail_start: false, - fail_request: false, - }) - .input(vec![ - Item::text(ItemKind::User, "initial").with_id("input-id"), - ]) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("prompts")).await.unwrap(); - finished(driver.next().await.unwrap()); - let requests = model.requests.lock().unwrap(); - assert_eq!(text(&requests[0].transcript[2]), "inference prompt"); - assert_eq!( - requests[0].transcript[2].id, - Some(agentkit_core::MessageId::new("input-id")) - ); - assert_eq!(text(&requests[0].transcript[0]), "inference-only context"); - let transcript = driver.snapshot().transcript; - assert_eq!(text(&transcript[1]), "logical prompt"); - assert_eq!(text(&transcript[0]), "persistent context"); - assert!(!transcript.iter().any(|item| { - item.parts - .iter() - .any(|part| matches!(part, Part::Text(text) if text.text.contains("inference"))) - })); -} - -#[tokio::test] -async fn failed_logical_prompt_edits_never_commit_and_request_edits_never_persist() { - for fail_start in [false, true] { - let model = ProbeModel::new([Script::Reply]); - let agent = Agent::builder() - .model(model.clone()) - .mutator(PromptHook { - fail_start, - fail_request: !fail_start, - }) - .input(vec![Item::text(ItemKind::User, "initial")]) - .build() - .unwrap(); - let mut driver = agent - .start(SessionConfig::new("prompt-failure")) - .await - .unwrap(); - assert!(matches!(driver.next().await, Err(LoopError::Mutator(_)))); - let transcript = driver.snapshot().transcript; - if fail_start { - assert!(transcript.is_empty()); - } else { - assert_eq!(transcript.len(), 2); - assert_eq!(text(&transcript[1]), "logical prompt"); - } - assert!(model.requests.lock().unwrap().is_empty()); - assert!(matches!( - driver.next().await.unwrap(), - LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) - )); - } -} - -#[derive(Default)] -struct StoredTranscript { - items: Mutex>, - appends: Mutex>, - rewrites: std::sync::atomic::AtomicUsize, -} -impl TranscriptObserver for Arc { - fn on_transcript_event(&self, event: TranscriptEvent<'_>) { - self.items.lock().unwrap().push(event.item.clone()); - self.appends.lock().unwrap().push(event.item.clone()); - } - fn on_transcript_rewrite(&self, event: TranscriptRewriteEvent<'_>) { - *self.items.lock().unwrap() = event.items.to_vec(); - self.rewrites.fetch_add(1, Ordering::SeqCst); - } -} - -#[tokio::test] -async fn second_turn_prompt_insertion_emits_rewrite_not_historic_appends() { - let stored = Arc::new(StoredTranscript::default()); - let agent = Agent::builder() - .model(ProbeModel::new([Script::Reply, Script::Reply])) - .mutator(PromptHook { - fail_start: false, - fail_request: false, - }) - .transcript_observer(stored.clone()) - .input(vec![Item::text(ItemKind::User, "first")]) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("rewrite")).await.unwrap(); - finished(driver.next().await.unwrap()); - assert_eq!(*stored.items.lock().unwrap(), driver.snapshot().transcript); - let before = stored.appends.lock().unwrap().len(); - driver - .submit_input_async(vec![Item::text(ItemKind::User, "second")]) - .await - .unwrap(); - finished(driver.next().await.unwrap()); - assert_eq!(*stored.items.lock().unwrap(), driver.snapshot().transcript); - assert_eq!(stored.rewrites.load(Ordering::SeqCst), 1); - assert_eq!( - stored.appends.lock().unwrap().len(), - before + 1, - "only the final output was appended; history/input arrived in the rewrite" - ); -} - -struct RejectInput; -#[async_trait] -impl LoopMutator for RejectInput { - async fn on_input(&self, input: &mut Vec, _: LifecycleCtx<'_>) -> Result<(), LoopError> { - if input.iter().any(|item| text(item) == "reject") { - Err(LoopError::Mutator("rejected batch".into())) - } else { - Ok(()) - } - } -} - -#[tokio::test] -async fn mixed_submission_rejection_preserves_earlier_valid_batch_and_order() { - let model = ProbeModel::new([Script::Reply]); - let agent = Agent::builder() - .model(model.clone()) - .mutator(RejectInput) - .build() - .unwrap(); - let mut driver = agent - .start(SessionConfig::new("mixed-input")) - .await - .unwrap(); - driver - .submit_input(vec![Item::text(ItemKind::User, "A")]) - .unwrap(); - assert!( - driver - .submit_input_async(vec![Item::text(ItemKind::User, "reject")]) - .await - .is_err() - ); - assert_eq!( - driver - .snapshot() - .pending_input - .iter() - .map(text) - .collect::>(), - ["A"] - ); - driver - .submit_input(vec![Item::text(ItemKind::User, "B")]) - .unwrap(); - driver - .submit_input_async(vec![Item::text(ItemKind::User, "C")]) - .await - .unwrap(); - assert_eq!( - driver - .snapshot() - .pending_input - .iter() - .map(text) - .collect::>(), - ["A", "B", "C"] - ); - finished(driver.next().await.unwrap()); - assert_eq!( - model.requests.lock().unwrap()[0] - .transcript - .iter() - .map(text) - .collect::>(), - ["A", "B", "C"] - ); -} - -#[tokio::test] -async fn queued_rejection_preserves_prior_accepted_and_later_unrelated_batches() { - let model = ProbeModel::new([Script::Reply]); - let agent = Agent::builder() - .model(model.clone()) - .mutator(RejectInput) - .build() - .unwrap(); - let mut driver = agent - .start(SessionConfig::new("queued-input")) - .await - .unwrap(); - for value in ["A", "reject", "C"] { - driver - .submit_input(vec![Item::text(ItemKind::User, value)]) - .unwrap(); - } - assert!(driver.next().await.is_err()); - assert_eq!( - driver - .snapshot() - .pending_input - .iter() - .map(text) - .collect::>(), - ["A", "C"] - ); - finished(driver.next().await.unwrap()); - assert_eq!( - model.requests.lock().unwrap()[0] - .transcript - .iter() - .map(text) - .collect::>(), - ["A", "C"] - ); -} - -#[derive(Clone)] -struct Gate { - open: Arc, - attempts: Arc, - completed: Arc, -} -impl Gate { - fn new() -> Self { - Self { - open: Arc::new(AtomicBool::new(false)), - attempts: Arc::default(), - completed: Arc::default(), - } - } - async fn wait(&self) { - self.attempts.fetch_add(1, Ordering::SeqCst); - if !self.open.load(Ordering::SeqCst) { - std::future::pending::<()>().await; - } - self.completed.fetch_add(1, Ordering::SeqCst); - } -} -struct BlockingHook { - phase: &'static str, - gate: Gate, -} -#[async_trait] -impl LoopMutator for BlockingHook { - async fn on_turn_finish( - &self, - _: TurnFinishPayload<'_>, - _: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - if self.phase == "finish" { - self.gate.wait().await; - } - Ok(()) - } - async fn on_turn_end(&self, _: &TurnResult, _: LifecycleCtx<'_>) -> Result<(), LoopError> { - if self.phase == "end" { - self.gate.wait().await; - } - Ok(()) - } - async fn on_session_end( - &self, - _: SessionEndPayload<'_>, - _: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - if self.phase == "session" { - self.gate.wait().await; - } - Ok(()) - } - async fn on_tool_batch( - &self, - _: ToolBatchPayload<'_>, - _: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - if self.phase == "batch" { - self.gate.wait().await; - } - Ok(()) - } -} -struct CompactCompletedTools; -#[async_trait] -impl LoopMutator for CompactCompletedTools { - async fn mutate( - &self, - cursor: &mut TranscriptCursor<'_>, - ctx: LoopCtx<'_>, - ) -> Result<(), LoopError> { - if ctx.point == MutationPoint::AfterToolResult { - cursor.retain(|item| { - !item - .parts - .iter() - .any(|part| matches!(part, Part::ToolCall(_) | Part::ToolResult(_))) - }); - } - Ok(()) - } -} - -#[tokio::test] -async fn batch_facts_and_completed_callbacks_survive_compaction_and_dropped_delivery() { - let events = Arc::default(); - let gate = Gate::new(); - let stored = Arc::new(StoredTranscript::default()); - let agent = Agent::builder() - .model(ProbeModel::new([Script::Tool, Script::Reply])) - .tool_executor(ProbeExecutor) - .mutator(ProbeHook::new("first", &events)) - .mutator(BlockingHook { - phase: "batch", - gate: gate.clone(), - }) - .mutator(ProbeHook::new("last", &events)) - .mutator(CompactCompletedTools) - .transcript_observer(stored.clone()) - .input(vec![Item::text(ItemKind::User, "input")]) - .build() - .unwrap(); - let mut driver = agent - .start(SessionConfig::new("captured-batch")) - .await - .unwrap(); - assert!( - tokio::time::timeout(std::time::Duration::from_millis(10), driver.next()) - .await - .is_err() - ); - assert_eq!(gate.attempts.load(Ordering::SeqCst), 1); - gate.open.store(true, Ordering::SeqCst); - finished(driver.next().await.unwrap()); - assert_eq!(gate.completed.load(Ordering::SeqCst), 1); - assert_eq!( - count(&events, "batch"), - 2, - "each completed first/last callback runs once" - ); - assert!( - !driver - .snapshot() - .transcript - .iter() - .flat_map(|item| &item.parts) - .any(|part| matches!(part, Part::ToolCall(_) | Part::ToolResult(_))) - ); - assert_eq!(*stored.items.lock().unwrap(), driver.snapshot().transcript); - driver.close().await.unwrap(); - assert_eq!(count(&events, "batch"), 2); -} - -struct CleanupManager { - inner: SimpleTaskManager, - gate: Gate, -} -#[async_trait] -impl TaskManager for CleanupManager { - async fn start_task( - &self, - request: TaskLaunchRequest, - ctx: TaskStartContext, - ) -> Result { - self.inner.start_task(request, ctx).await - } - async fn wait_for_turn( - &self, - id: &agentkit_core::TurnId, - cancel: Option, - ) -> Result, agentkit_task_manager::TaskManagerError> { - self.inner.wait_for_turn(id, cancel).await - } - async fn take_pending_loop_updates( - &self, - ) -> Result { - self.inner.take_pending_loop_updates().await - } - async fn on_turn_interrupted( - &self, - id: &agentkit_core::TurnId, - ) -> Result<(), agentkit_task_manager::TaskManagerError> { - self.gate.wait().await; - self.inner.on_turn_interrupted(id).await - } - fn handle(&self) -> agentkit_task_manager::TaskManagerHandle { - self.inner.handle() - } -} - -#[tokio::test] -async fn dropped_close_resumes_cleanup_and_terminal_callbacks_without_duplicates() { - for phase in ["cleanup", "finish", "end", "session"] { - let events = Arc::default(); - let observed: Arc>> = Arc::default(); - let gate = Gate::new(); - let cleanup_gate = Gate::new(); - if phase != "cleanup" { - cleanup_gate.open.store(true, Ordering::SeqCst); - } - let agent = Agent::builder() - .model(ProbeModel::new([Script::Tool])) - .tool_executor(ProbeExecutor) - .task_manager(CleanupManager { - inner: SimpleTaskManager::new(), - gate: cleanup_gate.clone(), - }) - .mutator(ProbeHook::new("first", &events)) - .mutator(BlockingHook { - phase, - gate: gate.clone(), - }) - .mutator(ProbeHook::new("last", &events)) - .observer(Capture(observed.clone())) - .input(vec![Item::text(ItemKind::User, "input")]) - .build() - .unwrap(); - let mut driver = agent - .start(SessionConfig::new("resumable-close")) - .await - .unwrap(); - assert!(matches!( - driver.next().await.unwrap(), - LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) - )); - assert!( - tokio::time::timeout(std::time::Duration::from_millis(10), driver.close()) - .await - .is_err(), - "{phase}" - ); - assert!(matches!( - driver.next().await, - Err(LoopError::InvalidState(_)) - )); - assert!( - driver - .submit_input(vec![Item::text(ItemKind::User, "closed")]) - .is_err() - ); - gate.open.store(true, Ordering::SeqCst); - cleanup_gate.open.store(true, Ordering::SeqCst); - driver.close().await.unwrap(); - driver.close().await.unwrap(); - assert_eq!(cleanup_gate.completed.load(Ordering::SeqCst), 1, "{phase}"); - if phase != "cleanup" { - assert_eq!(gate.completed.load(Ordering::SeqCst), 1, "{phase}"); - } - assert_eq!(count(&events, "turn_finish"), 2, "{phase}"); - assert_eq!(count(&events, "turn_end"), 2, "{phase}"); - assert_eq!(count(&events, "session_end"), 2, "{phase}"); - assert_eq!( - observed - .lock() - .unwrap() - .iter() - .filter(|event| matches!(event, AgentEvent::TurnFinished(_))) - .count(), - 1, - "{phase}" - ); - } -} - -struct InterruptProgress { - controller: agentkit_core::CancellationController, - terminals: Arc, -} -#[async_trait] -impl LoopMutator for InterruptProgress { - async fn on_model_progress( - &self, - _: &ModelTurnEvent, - _: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - self.controller.interrupt(); - Ok(()) - } - async fn on_turn_finish( - &self, - payload: TurnFinishPayload<'_>, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - assert_eq!(payload.result.finish_reason, FinishReason::Cancelled); - assert!(payload.output.is_none()); - assert!(ctx.cancellation.as_ref().unwrap().is_cancelled()); - self.terminals.fetch_add(1, Ordering::SeqCst); - Ok(()) - } - async fn on_turn_end( - &self, - result: &TurnResult, - ctx: LifecycleCtx<'_>, - ) -> Result<(), LoopError> { - assert_eq!(result.finish_reason, FinishReason::Cancelled); - assert!(ctx.cancellation.as_ref().unwrap().is_cancelled()); - self.terminals.fetch_add(1, Ordering::SeqCst); - Ok(()) - } -} -#[tokio::test] -async fn terminal_context_keeps_the_cancelled_active_checkpoint() { - let controller = agentkit_core::CancellationController::new(); - let terminals = Arc::default(); - let agent = Agent::builder() - .model(ProbeModel::new([Script::Reply])) - .cancellation(controller.handle()) - .mutator(InterruptProgress { - controller, - terminals: Arc::clone(&terminals), - }) - .input(vec![Item::text(ItemKind::User, "cancel during progress")]) - .build() - .unwrap(); - let mut driver = agent.start(SessionConfig::new("checkpoint")).await.unwrap(); - let result = finished(driver.next().await.unwrap()); - assert_eq!(result.finish_reason, FinishReason::Cancelled); - assert_eq!(terminals.load(Ordering::SeqCst), 2); - assert_eq!(driver.snapshot().transcript.last(), result.items.last()); -} - -#[tokio::test] -async fn readonly_end_notification_failure_preserves_committed_result() { - let events = Arc::default(); - let observed: Arc>> = Arc::default(); - let mut hook = ProbeHook::new("1", &events); - hook.fail = Some("turn_end"); - let agent = Agent::builder() - .model(ProbeModel::new([Script::Reply])) - .mutator(hook) - .observer(Capture(observed.clone())) - .input(vec![Item::text(ItemKind::User, "input")]) - .build() - .unwrap(); - let mut driver = agent - .start(SessionConfig::new("readonly-end")) - .await - .unwrap(); - let result = finished(driver.next().await.unwrap()); - assert_eq!(result.finish_reason, FinishReason::Completed); - assert_eq!(driver.snapshot().transcript.last(), result.items.last()); - assert!( - observed - .lock() - .unwrap() - .iter() - .any(|event| matches!(event, AgentEvent::TurnFinished(turn) if turn == &result)) - ); - assert_eq!(count(&events, "turn_end"), 1); - driver.close().await.unwrap(); - assert_eq!(count(&events, "turn_end"), 1); -} - -struct SendOnlyModel(ProbeModel); -struct SendOnlySession { - inner: ProbeSession, - calls: std::cell::Cell, -} -struct SendOnlyTurn { - inner: ProbeTurn, - polls: std::cell::Cell, -} -#[async_trait] -impl ModelAdapter for SendOnlyModel { - type Session = SendOnlySession; - async fn start_session(&self, config: SessionConfig) -> Result { - Ok(SendOnlySession { - inner: self.0.start_session(config).await?, - calls: std::cell::Cell::new(0), - }) - } -} -#[async_trait] -impl ModelSession for SendOnlySession { - type Turn = SendOnlyTurn; - async fn begin_turn( - &mut self, - request: TurnRequest, - cancel: Option, - ) -> Result { - self.calls.set(self.calls.get() + 1); - Ok(SendOnlyTurn { - inner: self.inner.begin_turn(request, cancel).await?, - polls: std::cell::Cell::new(0), - }) - } - fn provider_name(&self) -> Option<&str> { - Some("send-only-provider") - } - fn model_name(&self) -> Option<&str> { - Some("send-only-model") - } -} -#[async_trait] -impl ModelTurn for SendOnlyTurn { - async fn next_event( - &mut self, - cancel: Option, - ) -> Result, LoopError> { - self.polls.set(self.polls.get() + 1); - self.inner.next_event(cancel).await - } -} -fn require_send(value: T) -> T { - value -} - -#[tokio::test] -async fn public_lifecycle_futures_are_send_with_non_sync_session_and_turn() { - let events = Arc::default(); - let agent = Agent::builder() - .model(SendOnlyModel(ProbeModel::new([ - Script::Tool, - Script::Reply, - Script::BeginError, - ]))) - .tool_executor(ProbeExecutor) - .mutator(ProbeHook::new("send", &events)) - .build() - .unwrap(); - let mut driver = require_send(agent.start(SessionConfig::new("send-only"))) - .await - .unwrap(); - require_send(driver.submit_input_async(vec![Item::text(ItemKind::User, "tools")])) - .await - .unwrap(); - assert!(matches!( - require_send(driver.next()).await.unwrap(), - LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) - )); - finished(require_send(driver.next()).await.unwrap()); - require_send(driver.submit_input_async(vec![Item::text(ItemKind::User, "error")])) - .await - .unwrap(); - assert!(matches!( - require_send(driver.next()).await, - Err(LoopError::Provider(_)) - )); - require_send(driver.close()).await.unwrap(); -} - -#[tokio::test] -async fn earlier_queued_rejection_retains_new_async_input() { - let model = ProbeModel::new([Script::Reply]); - let agent = Agent::builder() - .model(model.clone()) - .mutator(RejectInput) - .build() - .unwrap(); - let mut driver = agent - .start(SessionConfig::new("inverse-mixed-input")) - .await - .unwrap(); - driver - .submit_input(vec![Item::text(ItemKind::User, "reject")]) - .unwrap(); - assert!( - driver - .submit_input_async(vec![Item::text(ItemKind::User, "B")]) - .await - .is_err() - ); - assert_eq!( - driver - .snapshot() - .pending_input - .iter() - .map(text) - .collect::>(), - ["B"] - ); - finished(driver.next().await.unwrap()); - assert_eq!( - model.requests.lock().unwrap()[0] - .transcript - .iter() - .map(text) - .collect::>(), - ["B"] - ); -} - -#[tokio::test] -async fn dropped_next_terminal_delivery_is_resumed_by_close() { - for phase in ["finish", "end"] { - let events = Arc::default(); - let gate = Gate::new(); - let observed: Arc>> = Arc::default(); - let agent = Agent::builder() - .model(ProbeModel::new([Script::Reply])) - .mutator(ProbeHook::new("first", &events)) - .mutator(BlockingHook { - phase, - gate: gate.clone(), - }) - .mutator(ProbeHook::new("last", &events)) - .observer(Capture(observed.clone())) - .input(vec![Item::text(ItemKind::User, "input")]) - .build() - .unwrap(); - let mut driver = agent - .start(SessionConfig::new("resume-next-terminal")) - .await - .unwrap(); - assert!( - tokio::time::timeout(std::time::Duration::from_millis(10), driver.next()) - .await - .is_err(), - "{phase}" - ); - gate.open.store(true, Ordering::SeqCst); - driver.close().await.unwrap(); - driver.close().await.unwrap(); - assert_eq!(count(&events, "turn_finish"), 2, "{phase}"); - assert_eq!(count(&events, "turn_end"), 2, "{phase}"); - assert_eq!(gate.completed.load(Ordering::SeqCst), 1, "{phase}"); - let observed = observed.lock().unwrap(); - let terminals: Vec<&TurnResult> = observed - .iter() - .filter_map(|event| match event { - AgentEvent::TurnFinished(result) => Some(result), - _ => None, - }) - .collect(); - assert_eq!(terminals.len(), 1, "{phase}"); - let result = terminals[0]; - assert_eq!( - result.finish_reason, - FinishReason::Completed, - "existing final outcome was resumed, not replaced" - ); - assert_eq!( - text(&result.items[0]), - "reply:mfirst:mlast:ffirst:flast", - "{phase}" - ); - assert_eq!( - driver.snapshot().transcript.last(), - result.items.last(), - "{phase}" - ); - } -} From 3d75d7179ad6ef4d675d6b9d88cc1864f59be759 Mon Sep 17 00:00:00 2001 From: daniel Date: Thu, 8 Oct 2026 20:36:35 +0100 Subject: [PATCH 3/3] feat(loop)!: add transactional mutation and awaited native delivery --- book/src/ch02-what-is-an-agent-loop.md | 2 +- book/src/ch06-driving-the-loop.md | 8 +- book/src/ch16-compaction.md | 16 +- book/src/ch19-reporting.md | 2 +- book/src/session-persistence.md | 33 +- crates/agentkit-loop/README.md | 22 +- crates/agentkit-loop/src/hook_tests.rs | 1245 ++++++++++++++++- crates/agentkit-loop/src/lib.rs | 834 +++++++++-- docs/architecture.md | 4 +- docs/compaction.md | 4 +- docs/loop.md | 2 +- .../src/main.rs | 53 +- 12 files changed, 2072 insertions(+), 153 deletions(-) diff --git a/book/src/ch02-what-is-an-agent-loop.md b/book/src/ch02-what-is-an-agent-loop.md index 98cd26e..01807f2 100644 --- a/book/src/ch02-what-is-an-agent-loop.md +++ b/book/src/ch02-what-is-an-agent-loop.md @@ -160,7 +160,7 @@ For loss-free transcript reconstruction (persistence, replication, audit), regis - **`LoopObserver`** sees a stream of `AgentEvent`s. Content arrives as deltas — partial fragments that don't carry their parent-`Item` identity — interleaved with lifecycle and telemetry events. Useful for UIs and logging, but a consumer cannot reassemble the canonical transcript from this stream alone. - **`TranscriptObserver`** fires exactly once per `Item` appended, receiving a session-addressed `TranscriptEvent` with the fully-formed `Item` ready to persist. Calls happen synchronously from the driver, in transcript order, at the single mutation point that owns the transcript — so what the observer sees is what the loop will send to the model on the next turn. -Mutator-driven rewrites (compaction, redaction, repair) do not fire `on_transcript_event`; they are signalled separately by `AgentEvent::MutationFinished`, which a persistence layer can use to snapshot the post-mutation state. +Mutator-driven rewrites (compaction, redaction, repair) replace history rather than appending to it, so they do not fire `on_transcript_event`. They fire the trait's other required method, `on_transcript_rewrite`, with the complete canonical transcript. Between the two, a persistence consumer cannot diverge from the driver. ## The three-layer model diff --git a/book/src/ch06-driving-the-loop.md b/book/src/ch06-driving-the-loop.md index 6d824ad..12efa7e 100644 --- a/book/src/ch06-driving-the-loop.md +++ b/book/src/ch06-driving-the-loop.md @@ -18,6 +18,7 @@ pub struct LoopDriver { mutators: Vec>, observers: Vec>, transcript_observers: Vec>, + deliveries: Vec>, transcript: Vec, pending_input: Vec, pending_approvals: BTreeMap, @@ -36,6 +37,7 @@ impl LoopDriver { -> Result<(), LoopError>; pub fn set_next_turn_cache(&mut self, cache: PromptCacheRequest) -> Result<(), LoopError>; pub fn snapshot(&self) -> LoopSnapshot; + pub fn take_delivery_errors(&mut self) -> DeliveryFailures; } ``` @@ -282,7 +284,9 @@ The full event taxonomy (`AgentEvent` is `#[non_exhaustive]` — keep a wildcard Observers are called inline, synchronously, in registration order. The loop task blocks briefly for each observer call. This is acceptable because observers should be fast — write to stderr, increment a counter, append to a buffer. Expensive processing should happen asynchronously behind a channel adapter. -For loss-free transcript reconstruction (persistence, replication, audit), the driver also fans out to a separate `TranscriptObserver` channel that fires once per `Item` appended — a session-addressed `TranscriptEvent` carrying the item, in transcript order. `LoopObserver` alone is not sufficient for this — content deltas span partial parts and historically tool results were appended without an event at all. Mutator-driven rewrites do **not** fire `on_transcript_event`; those are signaled by `AgentEvent::MutationFinished`. Register via `AgentBuilder::transcript_observer`. +For loss-free transcript reconstruction (persistence, replication, audit), the driver also fans out to a separate `TranscriptObserver` channel that fires once per `Item` appended — a session-addressed `TranscriptEvent` carrying the item, in transcript order. `LoopObserver` alone is not sufficient for this — content deltas span partial parts and historically tool results were appended without an event at all. Mutator-driven rewrites do **not** fire `on_transcript_event`; they fire the trait's other required method, `on_transcript_rewrite`, with the complete canonical transcript. Register via `AgentBuilder::transcript_observer`. + +Observers are synchronous and infallible, which is the right shape for telemetry but the wrong one for a host that has to *await* its own delivery of a fact at the moment it happens. That is what `AgentBuilder::delivery` is for: a `NativeDelivery` target is awaited at the emission site of each `NativeFact` — streaming progress as the model produces it, the read-only pre-commit `BeforeFinish`, and the post-commit `TurnFinished`. Delivery is read-only and cannot drive the loop, and a `DeliveryError` is diagnostics only: it never fails the operation, never rolls a commit back and never produces a second terminal event. Drain the diagnostics with `LoopDriver::take_delivery_errors`. ## Building the agent @@ -299,6 +303,7 @@ let agent = Agent::builder() .compaction(config) // default: none .observer(reporter) // default: none .transcript_observer(persistence) // default: none + .delivery(awaited_target) // default: none .transcript(vec![system_item]) // default: empty .input(vec![first_user_turn]) // default: empty (one-shot opener) .build()?; @@ -318,6 +323,7 @@ The builder validates that a model adapter is set. Everything else has sensible | `compaction` | `None` | Transcript grows without bounds | | `observers` | `[]` | No event reporting | | `transcript_observers` | `[]` | No transcript persistence hook | +| `deliveries` | `[]` | No awaited fact delivery | `Agent::start()` consumes the agent and returns a `LoopDriver` with the supplied transcript loaded passively. The first call to `next()` yields `AwaitingInput`; the host supplies the first user turn via `InputRequest::submit`, and the driver dispatches the model on the next `next()`. The agent's immutable configuration (adapter, tool sources, permissions) is moved into the driver. Multiple drivers can be created from the same `Agent` type by cloning it first. diff --git a/book/src/ch16-compaction.md b/book/src/ch16-compaction.md index 32e2599..fc729e9 100644 --- a/book/src/ch16-compaction.md +++ b/book/src/ch16-compaction.md @@ -261,7 +261,15 @@ This is why caching is configured separately from compaction in agentkit. Compac ## Loop integration -Compactors register as `LoopMutator`s. The loop runs every registered mutator at each `MutationPoint` — `AfterToolResult` (between tool results and the next inference call) and `AfterTurnEnded` (before the first inference call of a new turn, after pending input is appended). The trigger decides which points are relevant. +Compactors register as `LoopMutator`s. The loop runs every registered mutator at each `MutationPoint`: + +- `TurnStarted` — once when the driver creates a logical turn, after queued input is appended and before any other point or inference. This is the only point that also runs for turns that never dispatch an inference. +- `AfterToolResult` — between tool results and the next inference call. +- `AfterTurnEnded` — before the first inference call of a new turn. + +The trigger decides which points are relevant. A trigger that ignores `point` (such as `item_count_trigger`) now also gets the `TurnStarted` opportunity, so a transcript over the threshold compacts a little earlier in the turn. Filter on `point` if that matters. + +The chain is transactional. Mutators edit a candidate copy of the transcript; the loop validates it, re-checks cancellation, and only then assigns it to the live transcript in one synchronous step. A mutator that errors, produces a protocol-invalid transcript, cancels, or whose future is dropped leaves the live transcript exactly as it was. When a compactor fires: @@ -269,11 +277,15 @@ When a compactor fires: 2. The strategy pipeline transforms the transcript through the cursor 3. The loop validates transcript invariants (tool_use ↔ tool_result pairing) and hard-fails with `LoopError::Mutator` on a protocol violation 4. The compactor emits `AgentEvent::MutationFinished { mutator, dirty, metadata, .. }` +5. If the committed transcript actually differs from the live one, the loop delivers one `TranscriptObserver::on_transcript_rewrite` with the complete canonical transcript ```text Turn lifecycle with a registered compactor: - next() → merge pending input + next() → open turn, merge pending input + │ + ▼ + run mutators at TurnStarted │ ▼ begin model turn diff --git a/book/src/ch19-reporting.md b/book/src/ch19-reporting.md index fe7e4b3..7a32fa4 100644 --- a/book/src/ch19-reporting.md +++ b/book/src/ch19-reporting.md @@ -269,7 +269,7 @@ impl LoopObserver for AuditLogger { | Usage | `UsageUpdated` | | Diagnostic | `Warning` | -For loss-free transcript reconstruction, register a `TranscriptObserver` alongside `LoopObserver`. It fires a session-addressed `TranscriptEvent` once per `Item` appended, in transcript order — including the synthetic placeholder and the eventual real result for background-detached tools, correlated by `call_id` through the matching `ToolExecutionProgress` (placeholder) and `ToolResultReceived` (terminal result) events. +For loss-free transcript reconstruction, register a `TranscriptObserver` alongside `LoopObserver`. It fires a session-addressed `TranscriptEvent` once per `Item` appended, in transcript order — including the synthetic placeholder and the eventual real result for background-detached tools, correlated by `call_id` through the matching `ToolExecutionProgress` (placeholder) and `ToolResultReceived` (terminal result) events. Its other required method, `on_transcript_rewrite`, carries the complete canonical transcript when a mutator rewrote history instead of appending to it. ### Event timeline for a typical turn diff --git a/book/src/session-persistence.md b/book/src/session-persistence.md index cb04cdc..1c64239 100644 --- a/book/src/session-persistence.md +++ b/book/src/session-persistence.md @@ -6,21 +6,24 @@ The agent loop has no built-in storage backend. Persistence is intentionally a h | Primitive | Purpose | | ------------------------------------------------ | ----------------------------------------------------------------------------------------------------------------------- | -| `AgentBuilder::transcript(items)` | Restore prior transcript before the loop starts. | -| `TranscriptObserver::on_transcript_event(event)` | Mirror every newly-appended item to durable storage as the loop runs (`TranscriptEvent` carries `session_id` + `item`). | -| `LoopDriver::snapshot() -> LoopSnapshot` | Read-only point-in-time view of `transcript` and `pending_input` for ad-hoc dumps, audit, or full-state checkpoints. | +| `AgentBuilder::transcript(items)` | Restore prior transcript before the loop starts. | +| `TranscriptObserver::on_transcript_event(event)` | Mirror every newly-appended item to durable storage as the loop runs (`TranscriptEvent` carries `session_id` + `item`). | +| `TranscriptObserver::on_transcript_rewrite(event)` | Replace the stored transcript when a mutator rewrote history (`TranscriptRewriteEvent` carries `session_id` + `items`). | +| `LoopDriver::snapshot() -> LoopSnapshot` | Read-only point-in-time view of `transcript` and `pending_input` for ad-hoc dumps, audit, or full-state checkpoints. | That is the whole protocol. Any storage backend — in-memory map, sqlite, Postgres, S3, Redis — implements the same shape: 1. **On startup**: load the prior `Vec` for the session id (or empty for a fresh session) and pass it to `AgentBuilder::transcript`. -2. **During the run**: register a `TranscriptObserver` that appends each `Item` to durable storage. +2. **During the run**: register a `TranscriptObserver` that appends each `Item` to durable storage and replaces it on rewrite. 3. **On shutdown** (graceful or not): nothing required — the observer has already persisted every appended item. ## Two important guarantees **Append-only ordering.** `on_transcript_event` is called synchronously by the loop, in the exact order items land in the transcript. The observer is the single mutation point — every push to the transcript funnels through it. This means a strictly monotonic `seq` column on a sqlite `items` table reproduces the transcript byte-for-byte on reload. -**Mutators are out-of-band.** Mutator-driven transcript rewrites (compaction, redaction, repair) do **not** fire `on_transcript_event`. They are signalled via `AgentEvent::MutationFinished { dirty: true, .. }`, observable through a `LoopObserver`. A mutation-aware persistor subscribes to both channels and replaces the stored transcript when it sees a dirty mutation finish. An agent without mutators (most coding agents that rely on the provider's prompt cache plus a long context window) can ignore this. +**Mutator rewrites come through the second method.** Mutator-driven transcript rewrites (compaction, redaction, repair) replace history rather than appending to it, so they do **not** fire `on_transcript_event`. They fire `on_transcript_rewrite` with the complete canonical transcript — once per mutation point whose mutator chain produced a transcript that actually differs from the live one. Both methods are required, so an append-only consumer has to decide explicitly (a one-line no-op body) rather than silently diverging from the driver after a compaction pass. + +**Best-effort post-commit observation, not a commit gate.** Both methods are synchronous and infallible, and the loop has already committed by the time they run. A failed write cannot roll the turn back; a host that needs acknowledged durable publication owns that linearization point itself. ## A complete sqlite implementation @@ -72,6 +75,18 @@ let agent = Agent::builder() .build()?; ``` +```rust,ignore +impl TranscriptObserver for SqliteTranscriptObserver { + fn on_transcript_event(&self, event: TranscriptEvent<'_>) { + // append one item + } + + fn on_transcript_rewrite(&self, event: TranscriptRewriteEvent<'_>) { + // replace the stored rows with event.items + } +} +``` + That is the entire round-trip. Run the example twice with the same `--session` flag and the second run resumes mid-conversation — the first `next()` call returns `AwaitingInput` because the transcript is loaded but no input is queued, and the host supplies the next user message in response. ## Choosing a backend @@ -93,10 +108,10 @@ The integration test crate exercises the round-trip pattern internally; see `cra ## Mutation-aware persistence -If your agent registers any `LoopMutator`s (compaction, redaction, repair), the persistence flow extends: +If your agent registers any `LoopMutator`s (compaction, redaction, repair), the second observer method carries it: 1. `TranscriptObserver::on_transcript_event` continues to mirror new items as they arrive. -2. A `LoopObserver` subscribes to `AgentEvent::MutationFinished { dirty: true, .. }` and uses it as a signal to replace the stored transcript. -3. After a dirty mutation finish, call `LoopDriver::snapshot()` from the host's main task and replace the persisted transcript with `snapshot.transcript`. Subsequent `on_transcript_event` calls resume appending from the new tail. +2. `TranscriptObserver::on_transcript_rewrite` delivers the complete canonical transcript after a mutator chain committed a change. Replace the stored rows with `event.items`; subsequent `on_transcript_event` calls resume appending from the new tail. +3. A mutation that errored, produced a protocol-invalid transcript, was cancelled, or whose future was dropped leaves the live transcript untouched and notifies nothing — there is no partial rewrite to reconcile. A mutator that writes the same value back is not a change and notifies nothing either. -The two channels exist precisely so persistence can stay simple in the no-mutator case (one observer) without sacrificing correctness when mutators are wired in (one observer plus one event listener). +`AgentEvent::MutationStarted` / `MutationFinished` remain the mutator's own telemetry (which mutator ran, why, how much it replaced). They are not the persistence signal: a mutator chooses its own `dirty` label, while `on_transcript_rewrite` fires on the driver's own value comparison. diff --git a/crates/agentkit-loop/README.md b/crates/agentkit-loop/README.md index 62eeb64..b545967 100644 --- a/crates/agentkit-loop/README.md +++ b/crates/agentkit-loop/README.md @@ -25,10 +25,30 @@ Besides `mutate`, a `LoopMutator` registered with `AgentBuilder::mutator` can in - `on_session_start` edits session options before the model adapter starts the session. - `on_model_request` edits a single inference request (transcript, tools, cache, metadata). These edits are not persisted. -- `on_model_response` edits complete model output before it is committed, returned, or sent to tools. Content may change. Item identity, accounting and tool-call linkage may not. +- `on_model_response` edits complete model output before it is committed, returned, or sent to tools. Content may change — including tool-call arguments, which reach the executor. Item identity, accounting and tool-call linkage may not. `payload.disposition` says whether the loop will dispatch this response's tool calls (`ContinueWithTools`) or take the normal finish branch (`FinishTurnCandidate`), computed from the loop's own branch predicate rather than the response content. All hooks default to no-ops and run in registration order. Read-only notifications belong in `LoopObserver`; individual tool interception belongs at `ToolExecutor`. +## Transcript mutation + +`mutate` runs at every `MutationPoint`: `TurnStarted` (once per logical turn the driver creates, with queued input already appended, and the only point that also runs for turns that never dispatch an inference), `AfterToolResult`, and `AfterTurnEnded`. + +The chain is transactional. Mutators edit a candidate copy; the loop validates invariants, re-checks cancellation, then assigns the result to the live transcript in one synchronous step. A mutator that errors, produces a protocol-invalid transcript, cancels, or whose future is dropped leaves the live transcript untouched. A committed change is published to every `TranscriptObserver` as one `on_transcript_rewrite` carrying the complete canonical transcript; writing the same value back is not a change and publishes nothing. + +## Awaited delivery + +`LoopObserver` is synchronous and infallible. A host that must *await* its own delivery of a fact at the moment it happens registers a `NativeDelivery` with `AgentBuilder::delivery`. The driver awaits each target at the fact's emission site: + +- `NativeFact::Progress` — model deltas, usage, tool calls, attempt supersession, and the loop-authored background-detach placeholder, delivered as the driver consumes them rather than buffered until `next()` returns. +- `NativeFact::BeforeFinish` — the logical turn is about to finish, before its own terminal items are appended, so a consumer sees terminal output and cancellation partials before they commit. +- `NativeFact::TurnFinished` — the turn finished and its items are committed. + +The terminal pair is delivered exactly once per logical turn — including turns that end through cancellation, an error, a failed cleanup, or `retire_interrupted_turn` — for uninterrupted calls and for cooperative cancellation followed by retirement. A hard abort breaks that: `BeforeFinish` is awaited before anything commits, so dropping the `next()` future inside it loses the turn's terminal output candidate and leaves the turn active, and a later `retire_interrupted_turn` delivers a second prefinish for the same turn with a cancelled result. Treat prefinish as at-least-once if you drop driver futures. + +`HookCtx::cancellation` is `None` for the terminal facts, and also `None` for `Progress` when the agent was built without `AgentBuilder::cancellation` — match on the `NativeFact` variant rather than on the handle's presence. + +Delivery is read-only; a `DeliveryError` is diagnostics only — it never fails the operation, never stops the stream, never rolls a commit back, never replays and never produces a second terminal event, and later facts for the same turn still arrive. Drain it with `LoopDriver::take_delivery_errors`, which resets both the retained failures and the `dropped` count of what the bounded buffer discarded. + ## Quick start ```rust,no_run diff --git a/crates/agentkit-loop/src/hook_tests.rs b/crates/agentkit-loop/src/hook_tests.rs index f0b703a..79d3ac3 100644 --- a/crates/agentkit-loop/src/hook_tests.rs +++ b/crates/agentkit-loop/src/hook_tests.rs @@ -1,14 +1,22 @@ use super::*; +use agentkit_core::{Delta, PartId}; +use agentkit_task_manager::{AsyncTaskManager, RoutingDecision}; use agentkit_tools_core::{ToolContext, ToolExecutionOutcome, ToolResult}; use std::sync::{ Mutex, atomic::{AtomicBool, Ordering}, }; +use tokio::sync::Notify; +use tokio::time::{Duration, timeout}; #[derive(Clone, Copy)] enum Script { Reply, Tool, + /// `n` text deltas, then a plain reply. + Stream(usize), + /// A tool-call attempt the provider supersedes, then a plain reply. + Superseded, } #[derive(Clone, Default)] @@ -16,6 +24,9 @@ struct ProbeModel { scripts: Arc>>, requests: Arc>>, configs: Arc>>, + /// Awaited before every event after the first, so the model cannot run + /// ahead of an awaited delivery target. + gate: Option>, } impl ProbeModel { @@ -25,10 +36,19 @@ impl ProbeModel { ..Self::default() } } + + fn gated(mut self, gate: Arc) -> Self { + self.gate = Some(gate); + self + } } struct ProbeSession(ProbeModel); -struct ProbeTurn(VecDeque); +struct ProbeTurn { + events: VecDeque, + gate: Option>, + produced: usize, +} #[async_trait] impl ModelAdapter for ProbeModel { @@ -50,7 +70,7 @@ impl ModelSession for ProbeSession { self.0.requests.lock().unwrap().push(request); let script = self.0.scripts.lock().unwrap().pop_front(); let (item, finish_reason) = match script { - Some(Script::Reply) => ( + Some(Script::Reply | Script::Stream(_) | Script::Superseded) => ( Item::text(ItemKind::Assistant, "reply").with_id("message-reply"), FinishReason::Completed, ), @@ -69,6 +89,23 @@ impl ModelSession for ProbeSession { None => return Err(LoopError::Provider("unexpected inference".into())), }; let mut events = VecDeque::new(); + if let Some(Script::Stream(chunks)) = script { + let part_id = PartId::new("part-stream"); + for chunk in 0..chunks { + events.push_back(ModelTurnEvent::Delta(Delta::AppendText { + part_id: part_id.clone(), + chunk: format!("chunk-{chunk}"), + })); + } + } + if let Some(Script::Superseded) = script { + events.push_back(ModelTurnEvent::ToolCall(ToolCallPart::new( + "call-superseded", + "probe", + serde_json::json!({}), + ))); + events.push_back(ModelTurnEvent::ResponseAttemptSuperseded); + } if let Some(Part::ToolCall(call)) = item.parts.first() { events.push_back(ModelTurnEvent::ToolCall(call.clone())); } @@ -80,7 +117,11 @@ impl ModelSession for ProbeSession { model: None, response_id: None, })); - Ok(ProbeTurn(events)) + Ok(ProbeTurn { + events, + gate: self.0.gate.clone(), + produced: 0, + }) } } @@ -90,7 +131,13 @@ impl ModelTurn for ProbeTurn { &mut self, _: Option, ) -> Result, LoopError> { - Ok(self.0.pop_front()) + if let Some(gate) = &self.gate + && self.produced > 0 + { + gate.notified().await; + } + self.produced += 1; + Ok(self.events.pop_front()) } } @@ -456,3 +503,1193 @@ async fn cancelled_response_hook_finishes_cancelled_without_committing_output() .all(|part| !matches!(part, Part::ToolCall(_))) ); } + +// --------------------------------------------------------------------------- +// Transactional mutation, canonical rewrite, terminal funnel, delivery. +// +// Everything below drives the real `Agent` builder and `LoopDriver`, and reads +// results only through the public surface: `snapshot`, `next`, `submit`, +// `retire_interrupted_turn`, `take_delivery_errors` and the registered +// observer/delivery traits. +// --------------------------------------------------------------------------- + +/// One ordered notification, from whichever channel produced it. Transcript +/// notifications and delivered facts share a log so their relative order is +/// what the assertions read. +#[derive(Clone, Debug, PartialEq)] +enum Note { + Append(String), + Rewrite(Vec), + Progress(String), + BeforeFinish(FinishReason), + TurnFinished(FinishReason), +} + +#[derive(Clone, Default)] +struct Log(Arc>>); + +impl Log { + fn push(&self, note: Note) { + self.0.lock().unwrap().push(note); + } + + fn notes(&self) -> Vec { + self.0.lock().unwrap().clone() + } + + fn rewrites(&self) -> Vec> { + self.notes() + .into_iter() + .filter_map(|note| match note { + Note::Rewrite(items) => Some(items), + _ => None, + }) + .collect() + } + + fn count(&self, matching: fn(&Note) -> bool) -> usize { + self.notes().iter().filter(|note| matching(note)).count() + } +} + +fn label(item: &Item) -> String { + match item.parts.first() { + Some(Part::Text(text)) => text.text.clone(), + Some(Part::ToolCall(call)) => format!("call:{}", call.id.0), + Some(Part::ToolResult(result)) => format!("result:{}", result.call_id.0), + _ => format!("{:?}", item.kind), + } +} + +fn labels(items: &[Item]) -> Vec { + items.iter().map(label).collect() +} + +struct Persistor(Log); + +impl TranscriptObserver for Persistor { + fn on_transcript_event(&self, event: TranscriptEvent<'_>) { + self.0.push(Note::Append(label(event.item))); + } + + fn on_transcript_rewrite(&self, event: TranscriptRewriteEvent<'_>) { + self.0.push(Note::Rewrite(labels(event.items))); + } +} + +/// What a [`Mutation`] does to the candidate transcript. +#[derive(Clone, Copy)] +enum Edit { + /// Take `&mut` access and write the same value back. + Touch, + /// Append one item. + Append, + /// Replace the whole history with a single summary item. + Replace, + /// Append a `tool_result` that answers no call. + Orphan, + /// Remove everything, leaving nothing for the model to answer. + Erase, +} + +/// What a [`Mutation`] does after editing. +#[derive(Clone, Copy)] +enum After { + Return, + Fail, + Cancel, + /// Never returns, so the caller's future can be dropped mid-mutation. + Hang, +} + +/// Edits the candidate transcript at one mutation point, then returns, fails, +/// cancels or hangs — the four ways a mutation can end. +struct Mutation { + point: MutationPoint, + edit: Edit, + after: After, +} + +impl Mutation { + fn new(point: MutationPoint, edit: Edit, after: After) -> Self { + Self { point, edit, after } + } +} + +/// Records every mutation point it runs at, without editing anything. +struct RecordPoints(Arc>>); + +#[async_trait] +impl LoopMutator for RecordPoints { + async fn mutate( + &self, + _: &mut TranscriptCursor<'_>, + ctx: LoopCtx<'_>, + ) -> Result<(), LoopError> { + self.0.lock().unwrap().push(ctx.point); + Ok(()) + } +} + +/// Cancels once at [`MutationPoint::TurnStarted`], when the test arms it. +struct CancelWhenArmed(Arc); + +#[async_trait] +impl LoopMutator for CancelWhenArmed { + async fn mutate( + &self, + _: &mut TranscriptCursor<'_>, + ctx: LoopCtx<'_>, + ) -> Result<(), LoopError> { + if ctx.point == MutationPoint::TurnStarted && self.0.swap(false, Ordering::SeqCst) { + return Err(LoopError::Cancelled); + } + Ok(()) + } +} + +#[async_trait] +impl LoopMutator for Mutation { + async fn mutate( + &self, + cursor: &mut TranscriptCursor<'_>, + ctx: LoopCtx<'_>, + ) -> Result<(), LoopError> { + if ctx.point != self.point { + return Ok(()); + } + match self.edit { + Edit::Touch => { + if let Some(first) = cursor.first().cloned() { + cursor[0] = first; + } + } + Edit::Append => cursor.push(Item::text(ItemKind::Context, "injected")), + Edit::Replace => **cursor = vec![Item::text(ItemKind::User, "summary")], + Edit::Orphan => cursor.push(Item::new( + ItemKind::Tool, + vec![Part::ToolResult(ToolResultPart { + call_id: ToolCallId::new("orphan"), + output: ToolOutput::Text("orphan".into()), + is_error: false, + metadata: MetadataMap::new(), + })], + )), + Edit::Erase => cursor.clear(), + } + match self.after { + After::Return => Ok(()), + After::Fail => Err(LoopError::Mutator("mutation failed".into())), + After::Cancel => Err(LoopError::Cancelled), + After::Hang => std::future::pending().await, + } + } +} + +/// Records every delivered fact, optionally failing one kind of fact and +/// optionally releasing a model gate so progress timing is observable. +struct Target { + log: Log, + fail: Option, + gate: Option>, +} + +impl Target { + fn new(log: Log) -> Self { + Self { + log, + fail: None, + gate: None, + } + } + + fn failing(mut self, fact: NativeFactKind) -> Self { + self.fail = Some(fact); + self + } + + fn releasing(mut self, gate: Arc) -> Self { + self.gate = Some(gate); + self + } +} + +#[async_trait] +impl NativeDelivery for Target { + async fn deliver(&self, fact: NativeFact<'_>, ctx: HookCtx<'_>) -> Result<(), DeliveryError> { + assert!(ctx.turn_id.is_some(), "every fact belongs to a turn"); + match fact { + NativeFact::Progress(progress) => { + // A consumer sees `Progress` with `cancellation == None` when + // the agent was built without a cancellation handle, so the + // variant is the only thing that identifies a terminal fact. + let rendered = match progress { + NativeProgress::Model(ModelTurnEvent::Delta(Delta::AppendText { + chunk, + .. + })) => format!("delta:{chunk}"), + NativeProgress::Model(ModelTurnEvent::ToolCall(call)) => { + format!("call:{}", call.id.0) + } + NativeProgress::Model(ModelTurnEvent::ResponseAttemptSuperseded) => { + "superseded".into() + } + NativeProgress::Model(_) => "model".into(), + NativeProgress::ToolDetached(result) => { + format!("detached:{}", result.call_id.0) + } + }; + self.log.push(Note::Progress(rendered)); + if let Some(gate) = &self.gate { + gate.notify_one(); + } + } + NativeFact::BeforeFinish(result) => { + assert!( + ctx.cancellation.is_none(), + "terminal delivery runs uncancelled" + ); + self.log + .push(Note::BeforeFinish(result.finish_reason.clone())); + } + NativeFact::TurnFinished(result) => { + assert!( + ctx.cancellation.is_none(), + "terminal delivery runs uncancelled" + ); + self.log + .push(Note::TurnFinished(result.finish_reason.clone())); + } + } + let kind = match fact { + NativeFact::Progress(_) => NativeFactKind::Progress, + NativeFact::BeforeFinish(_) => NativeFactKind::BeforeFinish, + NativeFact::TurnFinished(_) => NativeFactKind::TurnFinished, + }; + match self.fail { + Some(failing) if failing == kind => Err(DeliveryError::new("target offline")), + _ => Ok(()), + } + } +} + +/// Records the disposition every response hook saw. +struct RecordDisposition(Arc>>); + +#[async_trait] +impl LoopMutator for RecordDisposition { + async fn on_model_response( + &self, + payload: ModelResponsePayload<'_>, + _: HookCtx<'_>, + ) -> Result<(), LoopError> { + self.0.lock().unwrap().push(payload.disposition); + Ok(()) + } +} + +#[tokio::test] +async fn failed_mutation_rolls_back_and_notifies_nothing() { + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .mutator(Mutation::new( + MutationPoint::TurnStarted, + Edit::Append, + After::Fail, + )) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let error = driver.next().await.unwrap_err(); + assert!(matches!(error, LoopError::Mutator(_)), "{error:?}"); + assert_eq!(labels(&driver.snapshot().transcript), ["hi"]); + assert!(log.rewrites().is_empty(), "{:?}", log.notes()); + assert!( + model.requests.lock().unwrap().is_empty(), + "a rolled-back mutation never reaches inference" + ); +} + +#[tokio::test] +async fn invalid_mutation_rolls_back_and_notifies_nothing() { + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .mutator(Mutation::new( + MutationPoint::TurnStarted, + Edit::Orphan, + After::Return, + )) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let error = driver.next().await.unwrap_err(); + assert!(matches!(error, LoopError::Mutator(_)), "{error:?}"); + assert_eq!(labels(&driver.snapshot().transcript), ["hi"]); + assert!(log.rewrites().is_empty(), "{:?}", log.notes()); +} + +#[tokio::test] +async fn dropped_mutation_rolls_back_and_notifies_nothing() { + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .mutator(Mutation::new( + MutationPoint::TurnStarted, + Edit::Replace, + After::Hang, + )) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + // Dropping the `next()` future abandons the candidate mid-chain. + assert!( + timeout(Duration::from_millis(50), driver.next()) + .await + .is_err(), + "the hanging mutation should not let next() resolve" + ); + assert_eq!(labels(&driver.snapshot().transcript), ["hi"]); + assert!(log.rewrites().is_empty(), "{:?}", log.notes()); +} + +#[tokio::test] +async fn cancelled_mutation_rolls_back_and_finishes_the_turn_cancelled() { + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone())) + .mutator(Mutation::new( + MutationPoint::TurnStarted, + Edit::Replace, + After::Cancel, + )) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("expected a finished turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Cancelled); + assert_eq!(labels(&driver.snapshot().transcript), ["hi"]); + assert!(log.rewrites().is_empty(), "{:?}", log.notes()); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 1, + "a cancelled turn start still ends exactly once: {:?}", + log.notes() + ); +} + +#[tokio::test] +async fn mutation_writing_the_same_value_notifies_nothing() { + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .mutator(Mutation::new( + MutationPoint::TurnStarted, + Edit::Touch, + After::Return, + )) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("expected a finished turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Completed); + assert!( + log.rewrites().is_empty(), + "taking &mut access is not a change: {:?}", + log.notes() + ); +} + +#[tokio::test] +async fn full_rewrite_notifies_one_canonical_snapshot() { + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .mutator(Mutation::new( + MutationPoint::TurnStarted, + Edit::Replace, + After::Return, + )) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Finished(_) = driver.next().await.unwrap() else { + panic!("expected a finished turn"); + }; + assert_eq!( + log.rewrites(), + [vec!["summary".to_owned()]], + "exactly one rewrite, carrying the whole canonical transcript" + ); + assert_eq!( + labels(&driver.snapshot().transcript), + ["summary", "reply"], + "the rewrite is what inference and later appends build on" + ); + let request = model.requests.lock().unwrap()[0].clone(); + assert_eq!(labels(&request.transcript), ["summary"]); +} + +#[tokio::test] +async fn turn_start_runs_once_for_a_turn_that_never_dispatches() { + let model = ProbeModel::new([]); + let log = Log::default(); + let points = Arc::new(Mutex::new(Vec::new())); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone())) + .mutator(RecordPoints(points.clone())) + .mutator(Mutation::new( + MutationPoint::AfterTurnEnded, + Edit::Erase, + After::Return, + )) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("expected a finished turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Completed); + assert!(result.items.is_empty()); + assert!( + model.requests.lock().unwrap().is_empty(), + "an emptied transcript must not dispatch an assistant prefill" + ); + assert_eq!( + points.lock().unwrap().as_slice(), + [MutationPoint::TurnStarted, MutationPoint::AfterTurnEnded], + "the turn-start point runs once, even though no inference followed" + ); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 1, + "{:?}", + log.notes() + ); +} + +#[tokio::test] +async fn turn_start_sees_admitted_input() { + let model = ProbeModel::new([Script::Reply]); + let seen = Arc::new(Mutex::new(Vec::>::new())); + + struct SeeInput(Arc>>>); + + #[async_trait] + impl LoopMutator for SeeInput { + async fn mutate( + &self, + cursor: &mut TranscriptCursor<'_>, + ctx: LoopCtx<'_>, + ) -> Result<(), LoopError> { + if ctx.point == MutationPoint::TurnStarted { + self.0.lock().unwrap().push(labels(cursor)); + } + Ok(()) + } + } + + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder.mutator(SeeInput(seen.clone())) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + let _ = driver.next().await.unwrap(); + + assert_eq!(seen.lock().unwrap().as_slice(), [vec!["hi".to_owned()]]); +} + +#[tokio::test] +async fn prefinish_precedes_the_terminal_commit_and_the_finish() { + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone())) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + let _ = driver.next().await.unwrap(); + + assert_eq!( + log.notes(), + [ + Note::Append("hi".into()), + Note::BeforeFinish(FinishReason::Completed), + Note::Append("reply".into()), + Note::TurnFinished(FinishReason::Completed), + ], + "prefinish observation lands before the terminal output commits" + ); +} + +#[tokio::test] +async fn cancelled_partial_is_delivered_before_it_commits() { + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone())) + .mutator(FailOnce::new(Hook::Request, || LoopError::Cancelled)) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("expected a finished turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Cancelled); + let notes = log.notes(); + let before = notes + .iter() + .position(|note| matches!(note, Note::BeforeFinish(_))) + .expect("prefinish delivered"); + let partial = notes + .iter() + .position(|note| matches!(note, Note::Append(text) if text.starts_with("Previous"))) + .expect("cancellation partial committed"); + assert!( + before < partial, + "the partial is observable before it commits: {notes:?}" + ); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 1, + "{notes:?}" + ); +} + +#[tokio::test] +async fn disposition_follows_the_loop_branch_through_tools_and_supersession() { + let model = ProbeModel::new([Script::Tool, Script::Superseded]); + let dispositions = Arc::new(Mutex::new(Vec::new())); + let executor = ProbeExecutor::default(); + let agent = agent(&model, &executor, |builder| { + builder.mutator(RecordDisposition(dispositions.clone())) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + loop { + match driver.next().await.unwrap() { + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => continue, + LoopStep::Finished(_) => break, + other => panic!("unexpected step {other:?}"), + } + } + + assert_eq!( + dispositions.lock().unwrap().as_slice(), + [ + ResponseDisposition::ContinueWithTools, + ResponseDisposition::FinishTurnCandidate, + ], + "the superseded attempt's tool call must not make the replacement \ + response look like a tool round" + ); +} + +#[tokio::test] +async fn disposition_cannot_be_moved_by_rewriting_linkage() { + let model = ProbeModel::new([Script::Tool]); + let executor = ProbeExecutor::default(); + let agent = agent(&model, &executor, |builder| builder.mutator(ForgeToolCall)); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let error = driver.next().await.unwrap_err(); + assert!(matches!(error, LoopError::Mutator(_)), "{error:?}"); + assert!(executor.0.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn progress_is_delivered_while_the_model_call_is_still_running() { + let gate = Arc::new(Notify::new()); + let model = ProbeModel::new([Script::Stream(3)]).gated(gate.clone()); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder.delivery(Target::new(log.clone()).releasing(gate.clone())) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + // The model refuses to produce its next event until the previous one has + // been delivered, so buffering progress until `next()` returns deadlocks. + let step = timeout(Duration::from_secs(5), driver.next()) + .await + .expect("progress must be delivered inside next(), not buffered") + .unwrap(); + assert!(matches!(step, LoopStep::Finished(_))); + assert_eq!( + log.notes() + .into_iter() + .take_while(|note| matches!(note, Note::Progress(_))) + .collect::>(), + [ + Note::Progress("delta:chunk-0".into()), + Note::Progress("delta:chunk-1".into()), + Note::Progress("delta:chunk-2".into()), + ] + ); +} + +#[tokio::test] +async fn a_tool_round_turn_ends_exactly_once_without_recommitting_output() { + let model = ProbeModel::new([Script::Tool, Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone())) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + loop { + match driver.next().await.unwrap() { + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => continue, + LoopStep::Finished(_) => break, + other => panic!("unexpected step {other:?}"), + } + } + + assert_eq!( + labels(&driver.snapshot().transcript), + ["hi", "call:call-probe", "result:call-probe", "reply"], + "the tool round's output commits once, for dispatch, not again at the finish" + ); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 1, + "{:?}", + log.notes() + ); + assert_eq!(log.count(|note| matches!(note, Note::BeforeFinish(_))), 1); +} + +#[tokio::test] +async fn an_errored_turn_ends_exactly_once() { + // No script left: the second inference fails inside the loop. + let model = ProbeModel::new([]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder.delivery(Target::new(log.clone())) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let error = driver.next().await.unwrap_err(); + assert!(matches!(error, LoopError::Provider(_)), "{error:?}"); + assert_eq!( + log.notes(), + [ + Note::BeforeFinish(FinishReason::Error), + Note::TurnFinished(FinishReason::Error), + ], + "an errored turn still ends once, through the same funnel" + ); +} + +/// Fails every foreground-task cleanup so retirement reports an error. +struct FailCleanup(SimpleTaskManager); + +#[async_trait] +impl TaskManager for FailCleanup { + async fn start_task( + &self, + request: agentkit_task_manager::TaskLaunchRequest, + ctx: agentkit_task_manager::TaskStartContext, + ) -> Result { + self.0.start_task(request, ctx).await + } + + async fn wait_for_turn( + &self, + turn_id: &agentkit_core::TurnId, + cancellation: Option, + ) -> Result, agentkit_task_manager::TaskManagerError> { + self.0.wait_for_turn(turn_id, cancellation).await + } + + async fn take_pending_loop_updates( + &self, + ) -> Result { + self.0.take_pending_loop_updates().await + } + + async fn on_turn_interrupted( + &self, + _: &agentkit_core::TurnId, + ) -> Result<(), agentkit_task_manager::TaskManagerError> { + Err(agentkit_task_manager::TaskManagerError::Internal( + "cleanup unavailable".into(), + )) + } + + fn handle(&self) -> agentkit_task_manager::TaskManagerHandle { + self.0.handle() + } +} + +/// Requires approval for every call, so a turn can be retired with tool work +/// still outstanding. +struct ApprovalExecutor; + +#[async_trait] +impl ToolExecutor for ApprovalExecutor { + fn specs(&self) -> Vec { + Vec::new() + } + async fn execute(&self, request: ToolRequest, _: &mut ToolContext<'_>) -> ToolExecutionOutcome { + ToolExecutionOutcome::Interrupted(agentkit_tools_core::ToolInterruption::ApprovalRequired( + ApprovalRequest { + task_id: None, + call_id: Some(request.call_id), + id: "approval:probe".into(), + request_kind: "probe".into(), + reason: agentkit_tools_core::ApprovalReason::SensitivePath, + summary: "probe needs approval".into(), + metadata: MetadataMap::new(), + }, + )) + } +} + +#[tokio::test] +async fn retirement_ends_the_turn_once_even_when_cleanup_fails() { + let model = ProbeModel::new([Script::Tool]); + let log = Log::default(); + let agent = Agent::builder() + .model(model.clone()) + .tool_executor(ApprovalExecutor) + .input(vec![Item::text(ItemKind::User, "hi")]) + .task_manager(FailCleanup(SimpleTaskManager::new())) + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone())) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_)) = driver.next().await.unwrap() + else { + panic!("expected an approval interrupt"); + }; + let error = driver.retire_interrupted_turn().await.unwrap_err(); + assert!(matches!(error, LoopError::Tool(_)), "{error:?}"); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 1, + "the ending is delivered before the cleanup error surfaces: {:?}", + log.notes() + ); + assert!( + driver.retire_interrupted_turn().await.unwrap().is_none(), + "a retired turn cannot be retired again" + ); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 1, + "{:?}", + log.notes() + ); +} + +#[tokio::test] +async fn delivery_failure_is_separate_from_the_operation_result() { + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone()).failing(NativeFactKind::TurnFinished)) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("a failed delivery must not fail the turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Completed); + assert_eq!(labels(&result.items), ["reply"]); + assert_eq!(labels(&driver.snapshot().transcript), ["hi", "reply"]); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 1, + "no replay and no second ending: {:?}", + log.notes() + ); + + let failures = driver.take_delivery_errors(); + assert_eq!(failures.failures.len(), 1, "{failures:?}"); + assert_eq!(failures.failures[0].fact, NativeFactKind::TurnFinished); + assert_eq!(failures.dropped, 0); + assert!( + driver.take_delivery_errors().is_empty(), + "draining clears the diagnostics" + ); +} + +#[tokio::test] +async fn failing_prefinish_delivery_leaves_the_turn_and_later_facts_intact() { + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone()).failing(NativeFactKind::BeforeFinish)) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("a failed prefinish delivery must not fail the turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Completed); + assert_eq!(labels(&result.items), ["reply"]); + assert_eq!( + log.notes(), + [ + Note::Append("hi".into()), + Note::BeforeFinish(FinishReason::Completed), + Note::Append("reply".into()), + Note::TurnFinished(FinishReason::Completed), + ], + "the terminal output still commits and the post-commit fact still \ + arrives after the prefinish target failed" + ); + assert_eq!(labels(&driver.snapshot().transcript), ["hi", "reply"]); + + let failures = driver.take_delivery_errors(); + assert_eq!(failures.failures.len(), 1, "{failures:?}"); + assert_eq!(failures.failures[0].fact, NativeFactKind::BeforeFinish); + assert_eq!(failures.dropped, 0); +} + +#[tokio::test] +async fn failing_progress_delivery_leaves_the_turn_and_later_facts_intact() { + let model = ProbeModel::new([Script::Stream(2)]); + let log = Log::default(); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone()).failing(NativeFactKind::Progress)) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("a failed progress delivery must not fail the turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Completed); + assert_eq!(labels(&result.items), ["reply"]); + assert_eq!( + log.notes(), + [ + Note::Append("hi".into()), + Note::Progress("delta:chunk-0".into()), + Note::Progress("delta:chunk-1".into()), + Note::BeforeFinish(FinishReason::Completed), + Note::Append("reply".into()), + Note::TurnFinished(FinishReason::Completed), + ], + "a failed progress delivery neither stops the stream nor the turn" + ); + + let failures = driver.take_delivery_errors(); + assert_eq!(failures.failures.len(), 2, "{failures:?}"); + assert!( + failures + .failures + .iter() + .all(|failure| failure.fact == NativeFactKind::Progress) + ); + assert_eq!(failures.dropped, 0); +} + +#[tokio::test] +async fn delivery_diagnostics_are_bounded_and_reset_on_drain() { + // More failures than the retained bound, so the buffer has to discard. + const FACTS: usize = 70; + let model = ProbeModel::new([Script::Stream(FACTS)]); + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder.delivery(Target::new(Log::default()).failing(NativeFactKind::Progress)) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + let LoopStep::Finished(_) = driver.next().await.unwrap() else { + panic!("expected a finished turn"); + }; + + let failures = driver.take_delivery_errors(); + assert!( + failures.failures.len() < FACTS, + "the buffer is bounded: {}", + failures.failures.len() + ); + assert_eq!( + failures.failures.len() + failures.dropped as usize, + FACTS, + "nothing vanishes from the accounting: {failures:?}" + ); + assert!( + failures + .failures + .iter() + .all(|failure| failure.fact == NativeFactKind::Progress) + ); + + assert!( + driver.take_delivery_errors().is_empty(), + "draining resets both the retained failures and the dropped count" + ); +} + +/// Routes every tool call to the background, so a turn ends by returning +/// `AwaitingInput` and the result arrives later as an out-of-band loop update. +fn background_task_manager() -> AsyncTaskManager { + AsyncTaskManager::new().routing(|_: &ToolRequest| RoutingDecision::Background) +} + +#[tokio::test] +async fn background_round_and_its_synthetic_turn_each_end_exactly_once() { + let model = ProbeModel::new([Script::Tool, Script::Reply]); + let log = Log::default(); + let points = Arc::new(Mutex::new(Vec::new())); + let agent = Agent::builder() + .model(model.clone()) + .tool_executor(ProbeExecutor::default()) + .input(vec![Item::text(ItemKind::User, "hi")]) + .task_manager(background_task_manager()) + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone())) + .mutator(RecordPoints(points.clone())) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + // The only call detached, so the turn ends through the `AwaitingInput` + // branch rather than a finished step. + let LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) = driver.next().await.unwrap() else { + panic!("expected the turn to end awaiting input"); + }; + assert_eq!( + log.count(|note| matches!(note, Note::BeforeFinish(_))), + 1, + "{:?}", + log.notes() + ); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 1, + "{:?}", + log.notes() + ); + assert_eq!( + log.count(|note| matches!(note, Note::Progress(text) if text.starts_with("detached:"))), + 1, + "the detach placeholder is delivered as progress: {:?}", + log.notes() + ); + + // The background result opens a synthetic turn that no host input asked + // for; it gets its own turn-start point and its own single ending. + timeout(Duration::from_secs(5), driver.wait_for_loop_update()) + .await + .expect("the background result should arrive") + .unwrap(); + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("expected the synthetic turn to finish"); + }; + assert_eq!(result.finish_reason, FinishReason::Completed); + assert_eq!( + points.lock().unwrap().as_slice(), + [ + MutationPoint::TurnStarted, + MutationPoint::AfterTurnEnded, + MutationPoint::TurnStarted, + MutationPoint::AfterTurnEnded, + ], + "the synthetic background-resolution turn runs the turn-start point too" + ); + assert_eq!( + log.count(|note| matches!(note, Note::BeforeFinish(_))), + 2, + "{:?}", + log.notes() + ); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 2, + "{:?}", + log.notes() + ); +} + +#[tokio::test] +async fn cancelling_a_synthetic_turn_start_requeues_its_background_resolutions() { + let model = ProbeModel::new([Script::Tool, Script::Reply]); + let armed = Arc::new(AtomicBool::new(false)); + let agent = Agent::builder() + .model(model.clone()) + .tool_executor(ProbeExecutor::default()) + .input(vec![Item::text(ItemKind::User, "hi")]) + .task_manager(background_task_manager()) + .mutator(CancelWhenArmed(armed.clone())) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) = driver.next().await.unwrap() else { + panic!("expected the turn to end awaiting input"); + }; + timeout(Duration::from_secs(5), driver.wait_for_loop_update()) + .await + .expect("the background result should arrive") + .unwrap(); + + // Cancel the synthetic turn at its start: it carries nothing yet, so its + // resolutions must survive for the next step rather than be dropped. + armed.store(true, Ordering::SeqCst); + let LoopStep::Finished(cancelled) = driver.next().await.unwrap() else { + panic!("expected a cancelled turn"); + }; + assert_eq!(cancelled.finish_reason, FinishReason::Cancelled); + let after_cancel = labels(&driver.snapshot().transcript); + assert_eq!( + after_cancel, + ["hi", "call:call-probe", "result:call-probe"], + "the cancelled turn start committed no resolution" + ); + + let LoopStep::Finished(result) = driver.next().await.unwrap() else { + panic!("expected the requeued resolution to drive a turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Completed); + assert_eq!(labels(&result.items), ["reply"]); + assert_eq!( + driver.snapshot().transcript.len(), + after_cancel.len() + 2, + "the requeued resolution and the reply both landed: {:?}", + labels(&driver.snapshot().transcript) + ); +} + +#[tokio::test] +async fn cancelling_pending_approvals_ends_the_turn_exactly_once() { + let model = ProbeModel::new([Script::Tool]); + let log = Log::default(); + let agent = Agent::builder() + .model(model.clone()) + .tool_executor(ApprovalExecutor) + .input(vec![Item::text(ItemKind::User, "hi")]) + .transcript_observer(Persistor(log.clone())) + .delivery(Target::new(log.clone())) + .build() + .unwrap(); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + let LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_)) = driver.next().await.unwrap() + else { + panic!("expected an approval interrupt"); + }; + let Some(LoopStep::Finished(result)) = driver.cancel_pending_approvals().await.unwrap() else { + panic!("expected cancelling the approvals to finish the turn"); + }; + assert_eq!(result.finish_reason, FinishReason::Cancelled); + assert_eq!( + log.count(|note| matches!(note, Note::BeforeFinish(_))), + 1, + "{:?}", + log.notes() + ); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 1, + "{:?}", + log.notes() + ); + + assert!( + driver.cancel_pending_approvals().await.unwrap().is_none(), + "there is nothing left to cancel" + ); + assert_eq!( + log.count(|note| matches!(note, Note::TurnFinished(_))), + 1, + "{:?}", + log.notes() + ); +} + +#[tokio::test] +async fn dropping_next_inside_prefinish_loses_the_output_and_retirement_redelivers() { + // The documented hard-abort hole: prefinish is awaited before anything + // commits, so a dropped future there strands the turn. + let model = ProbeModel::new([Script::Reply]); + let log = Log::default(); + + /// Hangs on prefinish so the test can drop the driver future inside it. + struct HangOnPrefinish(Log); + + #[async_trait] + impl NativeDelivery for HangOnPrefinish { + async fn deliver(&self, fact: NativeFact<'_>, _: HookCtx<'_>) -> Result<(), DeliveryError> { + if let NativeFact::BeforeFinish(result) = fact { + self.0 + .push(Note::BeforeFinish(result.finish_reason.clone())); + std::future::pending::<()>().await; + } + if let NativeFact::TurnFinished(result) = fact { + self.0 + .push(Note::TurnFinished(result.finish_reason.clone())); + } + Ok(()) + } + } + + let agent = agent(&model, &ProbeExecutor::default(), |builder| { + builder + .transcript_observer(Persistor(log.clone())) + .delivery(HangOnPrefinish(log.clone())) + }); + let mut driver = agent.start(SessionConfig::new("session")).await.unwrap(); + + assert!( + timeout(Duration::from_millis(50), driver.next()) + .await + .is_err(), + "the hanging prefinish delivery should not let next() resolve" + ); + assert_eq!( + labels(&driver.snapshot().transcript), + ["hi"], + "the terminal output candidate is lost, not committed" + ); + + // The turn is still active, so retirement re-runs the transition — and + // reaches the same hanging prefinish, which is why this is a timeout. + assert!( + timeout(Duration::from_millis(50), driver.retire_interrupted_turn()) + .await + .is_err(), + "retirement re-enters the hanging prefinish delivery" + ); + assert_eq!( + log.notes(), + [ + Note::Append("hi".into()), + Note::BeforeFinish(FinishReason::Completed), + Note::BeforeFinish(FinishReason::Cancelled), + ], + "prefinish is delivered twice for one logical turn: once for the lost \ + output candidate, once for the cancelled retirement" + ); +} diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index 4628613..0ebdd95 100644 --- a/crates/agentkit-loop/src/lib.rs +++ b/crates/agentkit-loop/src/lib.rs @@ -792,9 +792,11 @@ pub struct ObservedEvent { /// replication, or audit. /// /// Observers are called *synchronously* from inside the driver, in the -/// same order items land in the transcript. Compaction-driven transcript -/// rewrites do **not** fire `on_transcript_event` — those are signaled by -/// [`AgentEvent::CompactionFinished`] instead. +/// same order items land in the transcript. A [`LoopMutator`] that rewrites +/// the transcript does **not** produce appends; its committed result arrives +/// through [`TranscriptObserver::on_transcript_rewrite`]. Both methods are +/// required so a persistence consumer cannot silently diverge from the +/// driver after a compaction or redaction pass. /// /// Register via [`AgentBuilder::transcript_observer`]; multiple observers /// may be registered and are called in registration order. @@ -803,7 +805,7 @@ pub struct ObservedEvent { /// /// ```rust /// use agentkit_core::Item; -/// use agentkit_loop::{TranscriptEvent, TranscriptObserver}; +/// use agentkit_loop::{TranscriptEvent, TranscriptObserver, TranscriptRewriteEvent}; /// use std::sync::atomic::{AtomicUsize, Ordering}; /// /// struct CountingObserver { items: AtomicUsize } @@ -812,6 +814,10 @@ pub struct ObservedEvent { /// fn on_transcript_event(&self, _event: TranscriptEvent<'_>) { /// self.items.fetch_add(1, Ordering::Relaxed); /// } +/// +/// fn on_transcript_rewrite(&self, event: TranscriptRewriteEvent<'_>) { +/// self.items.store(event.items.len(), Ordering::Relaxed); +/// } /// } /// ``` pub trait TranscriptObserver: Send + Sync { @@ -820,6 +826,17 @@ pub trait TranscriptObserver: Send + Sync { /// state behind interior mutability so the driver can share an /// `Arc`. fn on_transcript_event(&self, event: TranscriptEvent<'_>); + + /// Replace the observer's transcript with this complete canonical + /// snapshot. Called synchronously once per mutation point whose mutator + /// chain committed a transcript that actually differs from the live one, + /// after invariant validation and the final cancellation check. + /// + /// This is post-commit, infallible, best-effort *observation*: it is not + /// a commit gate and it is not a durability contract. Hosts that need + /// acknowledged durable publication own that linearization point + /// themselves. Append-only consumers implement this as a no-op. + fn on_transcript_rewrite(&self, event: TranscriptRewriteEvent<'_>); } /// Session-addressed transcript append event delivered to @@ -832,10 +849,20 @@ pub struct TranscriptEvent<'a> { pub item: &'a Item, } +/// A complete committed transcript snapshot delivered to +/// [`TranscriptObserver`]s after a mutator chain rewrote the transcript. +#[derive(Clone, Debug)] +pub struct TranscriptRewriteEvent<'a> { + /// Session this rewrite belongs to. + pub session_id: &'a SessionId, + /// The full committed transcript, replacing whatever the observer held. + pub items: &'a [Item], +} + /// Where in the loop a [`LoopMutator`] is given a chance to modify the -/// transcript. Mutators run synchronously at these points; mid-stream -/// mutation (e.g. between content deltas) is intentionally not supported -/// because the assistant item is not yet fully constructed. +/// transcript. Mutators are awaited at these points; mid-stream mutation +/// (e.g. between content deltas) is intentionally not supported because the +/// assistant item is not yet fully constructed. #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] #[non_exhaustive] pub enum MutationPoint { @@ -845,6 +872,21 @@ pub enum MutationPoint { /// A new logical turn is about to make its first inference call. Pending /// user input has already been appended to the transcript. AfterTurnEnded, + /// The driver has just created a logical turn, once per turn, before any + /// other mutation point and before any inference for it. + /// + /// This is the only point that also covers turns that never dispatch an + /// inference — the synthetic turns the driver opens to carry background + /// tool resolutions or a late approval, and turns a mutator empties at + /// [`MutationPoint::AfterTurnEnded`]. Re-activating an existing turn (a + /// tool round that outlived its [`AgentEvent::TurnFinished`]) does not + /// run it again. + /// + /// Queued user input is already in the transcript, so a mutator can edit + /// the prompt the turn will answer. A synthetic turn's tool resolutions + /// are appended *after* this point, since the turn is created to carry + /// them; [`MutationPoint::AfterTurnEnded`] sees those. + TurnStarted, } /// Sink for emitting [`AgentEvent`]s from inside a [`LoopMutator`]. @@ -944,9 +986,29 @@ pub struct OutputItemPayload<'a> { pub metadata: &'a mut MetadataMap, } +/// What the loop will do with a complete model response once it is committed. +/// +/// Read-only, and computed from the loop's own branch predicate rather than +/// from the response content, so a hook does not have to re-derive loop +/// control flow. Tool-call linkage is immutable across the response hooks +/// (see [`ModelResponsePayload`]), so no edit can move the branch. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum ResponseDisposition { + /// The loop will dispatch this response's tool calls and the logical turn + /// continues. + ContinueWithTools, + /// This response takes the normal finish branch **if** the remaining hooks, + /// output validation and cancellation all succeed. It is a candidate, not + /// a committed outcome: an aborted operation commits nothing and finishes + /// no turn. + FinishTurnCandidate, +} + /// Complete model output before it is committed to the transcript, returned /// in the turn result, or dispatched to tools. Content may be rewritten; item /// identity, roles, accounting and tool-call linkage (ids and names) may not. +/// Tool-call *arguments* may be rewritten and reach the executor. /// Deltas already streamed to observers are not retracted. #[non_exhaustive] pub struct ModelResponsePayload<'a> { @@ -958,6 +1020,8 @@ pub struct ModelResponsePayload<'a> { pub usage: Option<&'a Usage>, /// Provider metadata for the call. pub metadata: &'a MetadataMap, + /// What the loop will do with this response once committed. + pub disposition: ResponseDisposition, } /// Async transcript mutator. Registered via [`AgentBuilder::mutator`] and @@ -1022,6 +1086,184 @@ pub trait LoopMutator: Send + Sync { } } +/// Awaited, read-only delivery of facts only the driver can observe. +/// +/// [`LoopObserver`] stays the synchronous telemetry fan-out. `NativeDelivery` +/// exists for the one thing a synchronous observer cannot do: let a host +/// *await* its own asynchronous delivery at the point the fact occurs. The +/// driver awaits each call at the fact's real emission site, so streaming +/// progress is delivered as the model produces it rather than buffered until +/// [`LoopDriver::next`] returns. +/// +/// Implementations are read-only. They cannot edit the transcript, the model +/// request, the model output or the turn result, and they cannot drive the +/// loop. Interception belongs to [`LoopMutator`]; individual tool +/// interception belongs at [`ToolExecutor`]. +/// +/// Failure is isolated: a [`DeliveryError`] never becomes a [`LoopError`], +/// never rolls back a commit, never re-runs interception or inference, and +/// never produces a second terminal event. Failures are recorded and drained +/// with [`LoopDriver::take_delivery_errors`]. +/// +/// Delivery of a fact that is already in flight cannot survive a hard abort: +/// dropping the future returned by [`LoopDriver::next`] loses whatever had +/// not been delivered, and the driver keeps no resumable delivery state. A +/// host that needs terminal facts delivered must let cooperative cancellation +/// and [`LoopDriver::retire_interrupted_turn`] run to completion. Every +/// exactly-once guarantee below is scoped to that: uninterrupted calls, or +/// cooperative cancellation followed by retirement. A dropped future breaks +/// it, and the driver does not pretend otherwise. +/// +/// [`HookCtx::cancellation`] is `None` for the terminal facts: terminal +/// delivery runs uncancelled so a read-only consumer cannot be starved of the +/// ending it needs to close its own state. For [`NativeFact::Progress`] it +/// carries the active turn's handle when the agent was built with +/// [`AgentBuilder::cancellation`], and is `None` otherwise — so `None` does +/// not by itself identify a terminal fact. Match on the [`NativeFact`] variant +/// for that. +/// +/// Register via [`AgentBuilder::delivery`]; multiple targets are awaited in +/// registration order. +#[async_trait] +pub trait NativeDelivery: Send + Sync { + /// Deliver one fact. See [`NativeFact`] for each variant's timing. + async fn deliver(&self, fact: NativeFact<'_>, ctx: HookCtx<'_>) -> Result<(), DeliveryError>; +} + +/// A loop fact handed to [`NativeDelivery`]. Payloads borrow the driver's own +/// values rather than introducing a parallel event taxonomy. +/// +/// Turn *start* is not here: [`MutationPoint::TurnStarted`] already delivers +/// it awaited, covers the same turns and can edit the prompt, so a read-only +/// duplicate would be strictly weaker. +#[derive(Clone, Copy, Debug)] +#[non_exhaustive] +pub enum NativeFact<'a> { + /// Non-terminal progress for the active turn, delivered as the driver + /// consumes it, before the matching [`AgentEvent`] reaches synchronous + /// observers. + Progress(NativeProgress<'a>), + /// The logical turn is about to finish. + /// + /// Read-only: the result is already final. Delivered before the turn's own + /// terminal items are appended to the transcript, so a consumer sees + /// cancellation partials and terminal model output before they commit. + /// Exactly once per logical turn, including turns that end through + /// cancellation, cleanup failure or [`LoopDriver::retire_interrupted_turn`]. + /// + /// Exactly-once holds for uninterrupted and cooperatively cancelled calls. + /// It does **not** survive a hard abort *inside this delivery*: because the + /// terminal items commit after it, dropping the [`LoopDriver::next`] future + /// while this fact is in flight loses that output candidate — it is never + /// committed and never returned — and leaves the turn active, so a later + /// [`LoopDriver::retire_interrupted_turn`] delivers a second prefinish for + /// the same logical turn, carrying a cancelled result instead. A consumer + /// that cannot tolerate that must not drop `next()`; one that chooses to + /// must treat prefinish as at-least-once and key on + /// [`TurnResult::turn_id`]. + BeforeFinish(&'a TurnResult), + /// The logical turn finished and its items are committed. Exactly once per + /// logical turn, paired with [`AgentEvent::TurnFinished`]. + /// + /// Unlike [`NativeFact::BeforeFinish`] this is the last step of the + /// transition, so the only way to miss it is to drop the driver future + /// before it runs. It is never delivered twice for one logical turn. + TurnFinished(&'a TurnResult), +} + +/// Non-terminal progress carried by [`NativeFact::Progress`]. +#[derive(Clone, Copy, Debug)] +#[non_exhaustive] +pub enum NativeProgress<'a> { + /// An event the model session produced for the active turn: deltas, usage, + /// tool calls and response-attempt supersession. [`ModelTurnEvent::Finished`] + /// is not progress — it reaches [`LoopMutator::on_model_response`] and then + /// [`NativeFact::BeforeFinish`]. + Model(&'a ModelTurnEvent), + /// The synthetic `tool_result` the driver authors when a tool call moves to + /// the background, paired with [`AgentEvent::ToolExecutionProgress`]. No + /// model produced it, so it does not arrive as a model event. + ToolDetached(&'a ToolResultPart), +} + +/// Failure reported by a [`NativeDelivery`] target. +/// +/// Deliberately not a [`LoopError`] variant: a read-only delivery failure must +/// never be mistaken for an operation failure. +#[derive(Clone, Debug, Error)] +#[error("native delivery failed: {message}")] +pub struct DeliveryError { + /// What went wrong, as reported by the delivery target. + pub message: String, +} + +impl DeliveryError { + /// Build a delivery failure from any message. + pub fn new(message: impl Into) -> Self { + Self { + message: message.into(), + } + } +} + +/// Which [`NativeFact`] a recorded delivery failure belongs to. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[non_exhaustive] +pub enum NativeFactKind { + /// [`NativeFact::Progress`]. + Progress, + /// [`NativeFact::BeforeFinish`]. + BeforeFinish, + /// [`NativeFact::TurnFinished`]. + TurnFinished, +} + +impl NativeFactKind { + fn of(fact: &NativeFact<'_>) -> Self { + match fact { + NativeFact::Progress(_) => Self::Progress, + NativeFact::BeforeFinish(_) => Self::BeforeFinish, + NativeFact::TurnFinished(_) => Self::TurnFinished, + } + } +} + +/// One recorded [`NativeDelivery`] failure. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct DeliveryFailure { + /// Session the fact belonged to. + pub session_id: SessionId, + /// Turn the fact belonged to, when one was active. + pub turn_id: Option, + /// Which fact failed to deliver. + pub fact: NativeFactKind, + /// The target's rendered error. + pub error: String, +} + +/// Drained [`NativeDelivery`] diagnostics. +/// +/// The driver keeps a bounded buffer so a permanently broken delivery target +/// cannot grow memory for the life of a session; `dropped` counts what the +/// bound discarded since the last drain. +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct DeliveryFailures { + /// Recorded failures, oldest first. + pub failures: Vec, + /// Failures discarded because the buffer was full. + pub dropped: u64, +} + +impl DeliveryFailures { + /// Returns `true` when nothing failed since the last drain. + pub fn is_empty(&self) -> bool { + self.failures.is_empty() && self.dropped == 0 + } +} + +/// Upper bound on retained [`DeliveryFailure`]s between drains. +const MAX_DELIVERY_FAILURES: usize = 64; + /// Lifecycle and streaming events emitted by the [`LoopDriver`]. /// /// Observers (see [`LoopObserver`]) receive these events in the order they @@ -1417,9 +1659,23 @@ struct ActiveToolRound { foreground_progressed: bool, } +/// Outcome of opening a logical turn. +enum OpenedTurn { + /// The turn is active and [`MutationPoint::TurnStarted`] has run. + Open(agentkit_core::TurnId), + /// A turn-start mutator cancelled, and the turn already finished through + /// the terminal funnel as this step. + Cancelled(Box), +} + #[derive(Default)] struct DriverLifecycle { active_turn: Option, + /// The active turn admitted queued input at [`MutationPoint::TurnStarted`] + /// and has not dispatched an inference for it yet. The input is already in + /// the transcript, so `pending_input` alone no longer answers "is there a + /// prompt waiting for the model?". + admitted_input: bool, } /// A configured agent ready to start a session. @@ -1477,6 +1733,7 @@ where mutators: Vec>, observers: Vec>, transcript_observers: Vec>, + deliveries: Vec>, transcript: Vec, input: Vec, telemetry: TelemetryConfig, @@ -1558,6 +1815,9 @@ where mutators: self.mutators.clone(), observers: self.observers.clone(), transcript_observers: self.transcript_observers.clone(), + deliveries: self.deliveries.clone(), + delivery_failures: VecDeque::new(), + dropped_delivery_failures: 0, transcript: self.transcript.clone(), pending_input: self.input.clone(), pending_approvals: BTreeMap::new(), @@ -1596,6 +1856,7 @@ where mutators: Vec>, observers: Vec>, transcript_observers: Vec>, + deliveries: Vec>, transcript: Vec, input: Vec, telemetry: TelemetryConfig, @@ -1617,6 +1878,7 @@ where mutators: Vec::new(), observers: Vec::new(), transcript_observers: Vec::new(), + deliveries: Vec::new(), transcript: Vec::new(), input: Vec::new(), telemetry: TelemetryConfig::default(), @@ -1723,6 +1985,19 @@ where self } + /// Register a [`NativeDelivery`] target that is awaited at the emission + /// site of each [`NativeFact`]. + /// + /// Multiple targets may be registered; they are awaited in registration + /// order. Use this only for facts a host must deliver *asynchronously* at + /// the moment they occur — streaming progress and the terminal + /// pre-commit/post-commit pair. Ordinary telemetry belongs in a + /// synchronous [`LoopObserver`]. + pub fn delivery(mut self, delivery: D) -> Self { + self.deliveries.push(Arc::new(delivery)); + self + } + /// Preload the driver's transcript with prior conversation state /// (defaults to empty). /// @@ -1777,6 +2052,7 @@ where mutators: self.mutators, observers: self.observers, transcript_observers: self.transcript_observers, + deliveries: self.deliveries, transcript: self.transcript, input: self.input, telemetry: self.telemetry, @@ -1833,6 +2109,13 @@ where mutators: Vec>, observers: Vec>, transcript_observers: Vec>, + deliveries: Vec>, + /// Bounded [`NativeDelivery`] diagnostics, oldest first, drained by + /// [`LoopDriver::take_delivery_errors`]. A read-only delivery failure is + /// never an operation failure, so it is recorded here instead of returned. + delivery_failures: VecDeque, + /// Failures discarded because `delivery_failures` was at capacity. + dropped_delivery_failures: u64, transcript: Vec, pending_input: Vec, pending_approvals: BTreeMap, @@ -1963,15 +2246,43 @@ where !self.pending_approvals.is_empty() } - fn start_logical_turn(&mut self) -> agentkit_core::TurnId { + /// Create a logical turn, admit queued input, and run + /// [`MutationPoint::TurnStarted`]. + /// + /// A no-op returning the active turn when one is already open, so every + /// caller can ask for "the turn this step belongs to" without tracking + /// whether it exists yet. + /// + /// `Err(LoopError::Cancelled)` means a turn-start mutator cancelled; the + /// caller must funnel the turn through [`Self::finish_cancelled`] rather + /// than leave it open. + async fn start_logical_turn(&mut self) -> Result { if let Some(turn_id) = &self.lifecycle.active_turn { - return turn_id.clone(); + return Ok(turn_id.clone()); } let turn_id = agentkit_core::TurnId::new(format!("turn-{}", self.next_turn_index)); self.next_turn_index += 1; - self.start_logical_turn_with(turn_id) + self.start_logical_turn_with(turn_id.clone()); + // Admit queued input before the mutation point so a prompt-editing + // mutator sees the turn's own input, and so turns that never reach + // an inference still get one chance to edit history. + let admitted: Vec = std::mem::take(&mut self.pending_input); + self.lifecycle.admitted_input = !admitted.is_empty(); + self.extend_transcript(admitted); + let cancellation = self + .cancellation + .as_ref() + .map(CancellationHandle::checkpoint); + self.run_mutators(MutationPoint::TurnStarted, Some(&turn_id), cancellation) + .await?; + Ok(turn_id) } + /// Activate `turn_id` without running [`MutationPoint::TurnStarted`]. + /// + /// Used to re-activate a turn whose tool round outlived its + /// [`AgentEvent::TurnFinished`], which is a continuation of work the + /// turn-start point already ran for — not a new turn. fn start_logical_turn_with(&mut self, turn_id: agentkit_core::TurnId) -> agentkit_core::TurnId { if let Some(active_turn) = &self.lifecycle.active_turn { return active_turn.clone(); @@ -1984,14 +2295,121 @@ where turn_id } - fn finish_logical_turn(&mut self, result: &TurnResult) { + /// [`Self::start_logical_turn`], funnelling a cancelling turn-start + /// mutator into the turn's cancelled ending instead of an error. + async fn open_logical_turn(&mut self) -> Result { + match self.start_logical_turn().await { + Ok(turn_id) => Ok(OpenedTurn::Open(turn_id)), + Err(LoopError::Cancelled) => { + let turn_id = self + .lifecycle + .active_turn + .clone() + .ok_or(LoopError::Cancelled)?; + self.finish_cancelled(turn_id, Vec::new()) + .await + .map(|step| OpenedTurn::Cancelled(Box::new(step))) + } + Err(error) => Err(error), + } + } + + /// The single terminal transition for a logical turn, for a turn whose + /// own terminal items are not in the transcript yet. + async fn finish_logical_turn(&mut self, result: TurnResult) -> TurnResult { + let items = result.items.clone(); + self.finish_logical_turn_inner(result, items).await + } + + /// [`Self::finish_logical_turn`] for a turn whose output was already + /// committed earlier in the turn so tools could be dispatched against it. + async fn finish_committed_logical_turn(&mut self, result: TurnResult) -> TurnResult { + self.finish_logical_turn_inner(result, Vec::new()).await + } + + /// Structural funnel every logical-turn ending passes through. + /// + /// The active-turn guard — not a per-path flag — is what makes the ending + /// happen exactly once: a second attempt for a turn that is no longer + /// active returns without delivering, emitting or appending anything. + /// + /// `uncommitted` are the turn's own terminal items that the transcript + /// does not hold yet (terminal model output, cancellation partials). They + /// are appended *after* [`NativeFact::BeforeFinish`] so a read-only + /// consumer observes them before they commit, and the turn never ends + /// owing a tool result. + /// + /// That ordering has one cost, and it is the deliberate trade: the awaited + /// prefinish delivery is a suspension point before anything has committed. + /// A caller that drops its future there loses `uncommitted` and leaves the + /// turn active, so [`Self::retire_interrupted_turn`] will run this + /// transition again — a second prefinish for the turn, with a cancelled + /// result. Buying exactly-once across that would mean resumable terminal + /// state, which this design rejects. + async fn finish_logical_turn_inner( + &mut self, + result: TurnResult, + uncommitted: Vec, + ) -> TurnResult { if self.pending_round_resume.as_ref() == Some(&result.turn_id) { self.pending_round_resume = None; } - if self.lifecycle.active_turn.as_ref() == Some(&result.turn_id) { - self.lifecycle.active_turn = None; - self.emit(AgentEvent::TurnFinished(result.clone())); + if self.lifecycle.active_turn.as_ref() != Some(&result.turn_id) { + return result; } + self.deliver(NativeFact::BeforeFinish(&result), None).await; + self.extend_transcript(uncommitted); + self.close_interrupted_tool_calls(); + self.lifecycle.active_turn = None; + self.lifecycle.admitted_input = false; + self.emit(AgentEvent::TurnFinished(result.clone())); + self.deliver(NativeFact::TurnFinished(&result), None).await; + result + } + + /// Await every registered [`NativeDelivery`] target for one fact. + /// + /// Infallible by construction: a target's [`DeliveryError`] is recorded as + /// diagnostics and the operation that produced the fact continues + /// unchanged — no rollback, no replay, no second terminal event. + async fn deliver(&mut self, fact: NativeFact<'_>, cancellation: Option) { + if self.deliveries.is_empty() { + return; + } + let kind = NativeFactKind::of(&fact); + let session_id = self.session_id.clone(); + // A terminal fact names its own turn: the active-turn slot is already + // cleared by the time `TurnFinished` is delivered. + let turn_id = match fact { + NativeFact::BeforeFinish(result) | NativeFact::TurnFinished(result) => { + Some(result.turn_id.clone()) + } + NativeFact::Progress(_) => self.lifecycle.active_turn.clone(), + }; + let deliveries = self.deliveries.clone(); + for target in &deliveries { + let ctx = HookCtx { + session_id: &session_id, + turn_id: turn_id.as_ref(), + cancellation: cancellation.clone(), + }; + if let Err(error) = target.deliver(fact, ctx).await { + self.record_delivery_failure(DeliveryFailure { + session_id: session_id.clone(), + turn_id: turn_id.clone(), + fact: kind, + error: error.to_string(), + }); + } + } + } + + fn record_delivery_failure(&mut self, failure: DeliveryFailure) { + if self.delivery_failures.len() >= MAX_DELIVERY_FAILURES { + self.delivery_failures.pop_front(); + self.dropped_delivery_failures = self.dropped_delivery_failures.saturating_add(1); + } + self.delivery_failures.push_back(failure); } fn emit_tool_catalog_events(&mut self, events: Vec) { @@ -2104,8 +2522,13 @@ where async fn drain_pending_loop_updates(&mut self) -> Result<(bool, Option), LoopError> { self.collect_pending_loop_updates().await?; let mut resolutions = std::mem::take(&mut self.pending_loop_updates); - if !resolutions.is_empty() { - self.start_logical_turn(); + if !resolutions.is_empty() + && let OpenedTurn::Cancelled(step) = self.open_logical_turn().await? + { + // The synthetic turn that would have carried these resolutions was + // cancelled at its start; requeue them for the next step. + self.pending_loop_updates = resolutions; + return Ok((false, Some(*step))); } let mut saw_items = false; while let Some(resolution) = resolutions.pop_front() { @@ -2115,7 +2538,7 @@ where saw_items = true; } TaskResolution::Approval(task) => { - let turn_id = self.start_logical_turn(); + let turn_id = self.start_logical_turn().await?; self.enqueue_pending_approval(&turn_id, task, None); } } @@ -2164,8 +2587,14 @@ where session_id: &observed_session_id, observers: &observers, }; + // The chain runs against a candidate copy, never the live history. + // Everything that can fail — a mutator, invariant validation, the + // final cancellation check — happens before the single synchronous + // assignment below, so a failed, invalid, cancelled or dropped chain + // leaves the live transcript exactly as it was and notifies nothing. + let mut candidate = self.transcript.clone(); let mut cursor = TranscriptCursor { - items: &mut self.transcript, + items: &mut candidate, dirty: false, }; for mutator in &mutators { @@ -2184,8 +2613,32 @@ where }; mutator.mutate(&mut cursor, ctx).await?; } - if cursor.dirty { - validate_transcript_invariants(cursor.items)?; + // `dirty` only proves a mutator took `&mut` access; it is a + // short-circuit, never evidence of change. A chain that writes the + // value back unchanged must commit nothing and notify nobody. + let dirty = cursor.dirty; + if !dirty || candidate == self.transcript { + return Ok(()); + } + validate_transcript_invariants(&candidate)?; + if cancellation + .as_ref() + .is_some_and(TurnCancellation::is_cancelled) + { + return Err(LoopError::Cancelled); + } + let now = Timestamp::now(); + for item in &mut candidate { + if item.created_at.is_none() { + item.created_at = Some(now); + } + } + self.transcript = candidate; + for observer in &self.transcript_observers { + observer.on_transcript_rewrite(TranscriptRewriteEvent { + session_id: &self.session_id, + items: &self.transcript, + }); } Ok(()) } @@ -2214,6 +2667,7 @@ where self.active_tool_round = None; return self .finish_cancelled(presentation_turn_id, Vec::new()) + .await .map(Some); } @@ -2264,7 +2718,8 @@ where TaskStartOutcome::Pending { kind, .. } => { self.emit(AgentEvent::ToolExecutionStarted(call.clone())); if kind == agentkit_task_manager::TaskKind::Background { - self.append_detach_placeholder(call.id.clone(), &call.name); + self.append_detach_placeholder(call.id.clone(), &call.name) + .await; if let Some(active) = self.active_tool_round.as_mut() { active.background_pending = true; } @@ -2299,7 +2754,8 @@ where } } Some(TurnTaskUpdate::Detached(snapshot)) => { - self.append_detach_placeholder(snapshot.call_id, &snapshot.tool_name); + self.append_detach_placeholder(snapshot.call_id, &snapshot.tool_name) + .await; if let Some(active) = self.active_tool_round.as_mut() { active.background_pending = true; active.foreground_progressed = true; @@ -2319,6 +2775,7 @@ where self.active_tool_round = None; return self .finish_cancelled(presentation_turn_id, Vec::new()) + .await .map(Some); } let active = self.active_tool_round.take().ok_or_else(|| { @@ -2383,7 +2840,9 @@ where { Ok(()) => {} Err(LoopError::Cancelled) => { - return self.finish_cancelled(turn_id, interrupted_assistant_items()); + return self + .finish_cancelled(turn_id, interrupted_assistant_items()) + .await; } Err(error) => return Err(error), } @@ -2394,14 +2853,15 @@ where // for the model to respond to. Finish the turn rather than dispatch an // assistant-prefill request, which most providers reject. if !transcript_has_pending_input(&self.transcript) { - let turn_result = TurnResult { - turn_id, - finish_reason: FinishReason::Completed, - items: Vec::new(), - usage: None, - metadata: MetadataMap::new(), - }; - self.finish_logical_turn(&turn_result); + let turn_result = self + .finish_logical_turn(TurnResult { + turn_id, + finish_reason: FinishReason::Completed, + items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + }) + .await; return Ok(LoopStep::Finished(turn_result)); } @@ -2409,7 +2869,9 @@ where .as_ref() .is_some_and(TurnCancellation::is_cancelled) { - return self.finish_cancelled(turn_id, interrupted_assistant_items()); + return self + .finish_cancelled(turn_id, interrupted_assistant_items()) + .await; } let catalog_events = self.tool_executor.drain_catalog_events(); @@ -2492,7 +2954,9 @@ where .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled(turn_id, interrupted_assistant_items()); + return self + .finish_cancelled(turn_id, interrupted_assistant_items()) + .await; } Err(error) => return Err(error), } @@ -2516,7 +2980,9 @@ where .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled(turn_id, interrupted_assistant_items()); + return self + .finish_cancelled(turn_id, interrupted_assistant_items()) + .await; } Err(error) => return Err(error), }; @@ -2557,20 +3023,31 @@ where .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled( - turn_id, - interrupted_stream_items(streamed_content.as_ref()), - ); + return self + .finish_cancelled(turn_id, interrupted_stream_items(streamed_content.as_ref())) + .await; } Err(error) => return Err(error), } { let attempt_superseded = matches!(event, ModelTurnEvent::ResponseAttemptSuperseded); if attempt_superseded { + // Drop the superseded attempt's reconstruction state before + // the awaited delivery below, so a cancellation observed + // during it cannot commit that attempt's partial content. saw_tool_call = false; latest_usage = None; if let Some(content) = &mut streamed_content { content.reset(); } + } + if !matches!(event, ModelTurnEvent::Finished(_)) { + self.deliver( + NativeFact::Progress(NativeProgress::Model(&event)), + cancellation.clone(), + ) + .await; + } + if attempt_superseded { self.emit(AgentEvent::ResponseAttemptSuperseded); } if cancellation @@ -2582,10 +3059,9 @@ where .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled( - turn_id, - interrupted_stream_items(streamed_content.as_ref()), - ); + return self + .finish_cancelled(turn_id, interrupted_stream_items(streamed_content.as_ref())) + .await; } if attempt_superseded { continue; @@ -2622,11 +3098,20 @@ where LoopError::Provider("model turn ended without a Finished event".into()) })?; result.usage = merge_usage(result.usage, latest_usage); + // Taken from the loop's own branch predicate, after every superseded + // attempt reset it, so a hook never has to re-derive control flow from + // the response content. + let disposition = if saw_tool_call { + ResponseDisposition::ContinueWithTools + } else { + ResponseDisposition::FinishTurnCandidate + }; match run_response_hooks( &self.mutators, &self.session_id, &turn_id, cancellation.as_ref(), + disposition, &mut result, ) .instrument(chat_span.clone()) @@ -2638,10 +3123,9 @@ where .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled( - turn_id, - interrupted_stream_items(streamed_content.as_ref()), - ); + return self + .finish_cancelled(turn_id, interrupted_stream_items(streamed_content.as_ref())) + .await; } Err(error) => return Err(error), } @@ -2705,9 +3189,11 @@ where item }) .collect(); - self.extend_transcript(output_items.clone()); - if saw_tool_call { + // Tool dispatch reads the committed transcript, so this output + // commits before the round rather than through the terminal + // funnel. + self.extend_transcript(output_items.clone()); let pending_calls = extract_tool_calls(&output_items) .into_iter() .map(|call| { @@ -2733,13 +3219,14 @@ where if let Some(step) = self.continue_active_tool_round().await? { return Ok(step); } - self.finish_logical_turn(&TurnResult { + self.finish_committed_logical_turn(TurnResult { turn_id, finish_reason: result.finish_reason, items: output_items, usage: result.usage, metadata: result.metadata, - }); + }) + .await; return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput( InputRequest { session_id: self.session_id.clone(), @@ -2748,14 +3235,15 @@ where ))); } - let turn_result = TurnResult { - turn_id, - finish_reason: result.finish_reason, - items: output_items, - usage: result.usage, - metadata: result.metadata, - }; - self.finish_logical_turn(&turn_result); + let turn_result = self + .finish_logical_turn(TurnResult { + turn_id, + finish_reason: result.finish_reason, + items: output_items, + usage: result.usage, + metadata: result.metadata, + }) + .await; Ok(LoopStep::Finished(turn_result)) } @@ -2844,7 +3332,8 @@ where self.append_detach_placeholder( pending.call.id.clone(), &pending.call.name, - ); + ) + .await; } else { self.active_tool_round = Some(ActiveToolRound { presentation_turn_id: pending.presentation_turn_id.clone(), @@ -2890,23 +3379,25 @@ where } } - fn finish_cancelled( + /// Finish `turn_id` as cancelled, carrying whatever partial output the + /// turn produced. The partial is read-only: it reaches the terminal facts + /// and the transcript, but no hook may rewrite it. + async fn finish_cancelled( &mut self, turn_id: agentkit_core::TurnId, items: Vec, ) -> Result { let pending = self.drain_pending_approval_items(); self.reject_drained_approvals(pending); - self.extend_transcript(items.clone()); - self.close_interrupted_tool_calls(); - let turn_result = TurnResult { - turn_id, - finish_reason: FinishReason::Cancelled, - items, - usage: None, - metadata: interrupted_metadata("turn"), - }; - self.finish_logical_turn(&turn_result); + let turn_result = self + .finish_logical_turn(TurnResult { + turn_id, + finish_reason: FinishReason::Cancelled, + items, + usage: None, + metadata: interrupted_metadata("turn"), + }) + .await; Ok(LoopStep::Finished(turn_result)) } @@ -3040,7 +3531,10 @@ where /// /// This clears the blocking approval and appends an error tool result so /// the transcript remains provider-valid if the host continues the turn. - pub fn cancel_pending_approval_for(&mut self, call_id: ToolCallId) -> Result<(), LoopError> { + pub async fn cancel_pending_approval_for( + &mut self, + call_id: ToolCallId, + ) -> Result<(), LoopError> { let Some(pending) = self.drain_pending_approval_for(&call_id) else { return Err(LoopError::InvalidState(format!( "no approval request is pending for call {}", @@ -3050,7 +3544,7 @@ where let turn_id = pending.presentation_turn_id.clone(); self.reject_drained_approvals(vec![pending]); if self.pending_approvals.is_empty() && self.active_tool_round.is_none() { - let _ = self.finish_cancelled(turn_id, Vec::new())?; + let _ = self.finish_cancelled(turn_id, Vec::new()).await?; } Ok(()) } @@ -3096,17 +3590,19 @@ where } self.reject_drained_approvals(pending); if let Some(error) = cleanup_error { - self.close_interrupted_tool_calls(); - self.finish_logical_turn(&TurnResult { + // Cleanup failed, but the turn still ends exactly once and the + // transcript still closes its tool calls — the funnel owns both. + self.finish_logical_turn(TurnResult { turn_id, finish_reason: FinishReason::Error, items: Vec::new(), usage: None, metadata: MetadataMap::new(), - }); + }) + .await; return Err(error); } - self.finish_cancelled(turn_id, Vec::new()).map(Some) + self.finish_cancelled(turn_id, Vec::new()).await.map(Some) } /// Retire the active logical turn without resuming model or tool execution. @@ -3120,6 +3616,13 @@ where /// Returns `None` when no logical turn is active, including on repeated calls. /// Otherwise emits exactly one cancelled [`AgentEvent::TurnFinished`]. /// + /// This is the designated recovery path after a dropped [`Self::next`] + /// future, and it is the one case where a turn can see a second + /// [`NativeFact::BeforeFinish`]: if the drop happened inside that + /// delivery, the turn is still active and its terminal output candidate was + /// never committed, so retirement re-runs the transition with a cancelled + /// result. See [`NativeFact::BeforeFinish`]. + /// /// # Errors /// /// Returns the first task cleanup error, after completing local cleanup and @@ -3133,18 +3636,34 @@ where // At AfterToolResult the foreground task round has already completed. self.pending_round_resume = None; let cleanup = self.cleanup_interrupted_turn().await; - let result = TurnResult { - turn_id, - finish_reason: FinishReason::Cancelled, - items: Vec::new(), - usage: None, - metadata: interrupted_metadata("turn"), - }; - self.finish_logical_turn(&result); + // The terminal facts are delivered before the cleanup error surfaces, + // so a failing task manager cannot withhold the turn's ending. + let result = self + .finish_logical_turn(TurnResult { + turn_id, + finish_reason: FinishReason::Cancelled, + items: Vec::new(), + usage: None, + metadata: interrupted_metadata("turn"), + }) + .await; cleanup?; Ok(Some(result)) } + /// Drain the [`NativeDelivery`] failures recorded since the last call. + /// + /// Delivery failures are diagnostics, never operation failures: the work + /// that produced the undelivered fact is already committed and its own + /// result was returned unchanged. Hosts that care about delivery drain + /// this after each driver call and surface it separately. + pub fn take_delivery_errors(&mut self) -> DeliveryFailures { + DeliveryFailures { + failures: std::mem::take(&mut self.delivery_failures).into(), + dropped: std::mem::take(&mut self.dropped_delivery_failures), + } + } + /// Take a read-only snapshot of the driver's current transcript and input queue. pub fn snapshot(&self) -> LoopSnapshot { LoopSnapshot { @@ -3196,40 +3715,26 @@ where /// Returns [`LoopError::InvalidState`] if called while an unresolved /// interrupt is pending, or propagates provider / tool / compaction errors. pub async fn next(&mut self) -> Result { - if self.lifecycle.active_turn.is_none() { - let continuation_turn = self - .pending_approval_order - .iter() - .find_map(|call_id| self.pending_approvals.get(call_id)) - .map(|pending| pending.presentation_turn_id.clone()) - .or_else(|| { - self.active_tool_round - .as_ref() - .map(|active| active.presentation_turn_id.clone()) - }) - .or_else(|| self.pending_round_resume.clone()); - if let Some(turn_id) = continuation_turn { - self.start_logical_turn_with(turn_id); - } else if !self.pending_input.is_empty() { - self.start_logical_turn(); - } - } - let result = self.next_inner().await; match &result { - Ok(LoopStep::Finished(turn)) => self.finish_logical_turn(turn), + Ok(LoopStep::Finished(turn)) => { + // Already funnelled by whichever path produced it; this is the + // guard's no-op for a turn that is no longer active. + self.finish_committed_logical_turn(turn.clone()).await; + } Err(_) => { if let Some(turn_id) = self.lifecycle.active_turn.clone() { if let Err(error) = self.cleanup_interrupted_turn().await { tracing::debug!(%error, "failed to clean up turn after loop error"); } - self.finish_logical_turn(&TurnResult { + self.finish_logical_turn(TurnResult { turn_id, finish_reason: FinishReason::Error, items: Vec::new(), usage: None, metadata: MetadataMap::new(), - }); + }) + .await; } } _ => {} @@ -3274,7 +3779,45 @@ where } } + /// Open the logical turn the next step belongs to, before any of that + /// step's work. + /// + /// Continuation work re-activates the turn it belongs to; otherwise queued + /// input opens a fresh one. Returns the cancelled terminal step when a + /// turn-start mutator cancels. + async fn open_step_turn(&mut self) -> Result, LoopError> { + if self.lifecycle.active_turn.is_some() { + return Ok(None); + } + let continuation_turn = self + .pending_approval_order + .iter() + .find_map(|call_id| self.pending_approvals.get(call_id)) + .map(|pending| pending.presentation_turn_id.clone()) + .or_else(|| { + self.active_tool_round + .as_ref() + .map(|active| active.presentation_turn_id.clone()) + }) + .or_else(|| self.pending_round_resume.clone()); + if let Some(turn_id) = continuation_turn { + self.start_logical_turn_with(turn_id); + return Ok(None); + } + if self.pending_input.is_empty() { + return Ok(None); + } + match self.open_logical_turn().await? { + OpenedTurn::Open(_) => Ok(None), + OpenedTurn::Cancelled(step) => Ok(Some(*step)), + } + } + async fn next_inner(&mut self) -> Result { + if let Some(step) = self.open_step_turn().await? { + return Ok(step); + } + if let Some(pending) = self.take_next_resolved_approval() { return self.resume_after_approval(pending).await; } @@ -3299,11 +3842,24 @@ where // before unrelated background completions so a delayed approval cannot // bind itself to that turn's TurnStarted event. AfterToolResult resumes // remain ordered ahead of fresh input below. - if self.pending_round_resume.is_none() && !self.pending_input.is_empty() { + // + // Two queues feed this branch: `admitted_input` for input this step + // already folded into the transcript when it opened the turn, and + // `pending_input` for input queued against a turn that was already + // open. Only `pending_round_resume` defers them, and the branch that + // takes it below drains `pending_input` itself — so admitted input is + // never left unanswered behind the `AwaitingInput` tail. + if self.pending_round_resume.is_none() + && (self.lifecycle.admitted_input || !self.pending_input.is_empty()) + { // Take updates now to preserve the driver's once-per-step manager // handoff, but defer presenting them until this input turn ends. self.collect_pending_loop_updates().await?; - let turn_id = self.start_logical_turn(); + let turn_id = match self.open_logical_turn().await? { + OpenedTurn::Open(turn_id) => turn_id, + OpenedTurn::Cancelled(step) => return Ok(*step), + }; + self.lifecycle.admitted_input = false; let drained: Vec = std::mem::take(&mut self.pending_input); self.extend_transcript(drained); return self @@ -3337,7 +3893,11 @@ where ))); } - let turn_id = self.start_logical_turn(); + let turn_id = match self.open_logical_turn().await? { + OpenedTurn::Open(turn_id) => turn_id, + OpenedTurn::Cancelled(step) => return Ok(*step), + }; + self.lifecycle.admitted_input = false; let drained: Vec = std::mem::take(&mut self.pending_input); self.extend_transcript(drained); self.drive_turn(turn_id, MutationPoint::AfterTurnEnded) @@ -3365,7 +3925,7 @@ where self.transcript.push(item); } - fn append_detach_placeholder(&mut self, call_id: ToolCallId, tool_name: &str) { + async fn append_detach_placeholder(&mut self, call_id: ToolCallId, tool_name: &str) { self.background_call_ids.insert(call_id.clone()); if !self.detached_call_ids.insert(call_id.clone()) { return; @@ -3378,6 +3938,12 @@ where is_error: false, metadata: MetadataMap::new(), }; + let cancellation = self.tool_cancellation_for(&call_id, None); + self.deliver( + NativeFact::Progress(NativeProgress::ToolDetached(&detached_result)), + cancellation, + ) + .await; self.emit(AgentEvent::ToolExecutionProgress(detached_result.clone())); self.append_item(Item { id: None, @@ -5149,6 +5715,7 @@ async fn run_response_hooks( session_id: &SessionId, turn_id: &agentkit_core::TurnId, cancellation: Option<&TurnCancellation>, + disposition: ResponseDisposition, result: &mut ModelTurnResult, ) -> Result<(), LoopError> { if mutators.is_empty() { @@ -5170,6 +5737,7 @@ async fn run_response_hooks( finish_reason: &result.finish_reason, usage: result.usage.as_ref(), metadata: &result.metadata, + disposition, }; let ctx = HookCtx { session_id, @@ -6466,6 +7034,22 @@ mod tests { } } + /// Records the loop-authored detach placeholders delivered as progress. + #[derive(Clone, Default)] + struct RecordingDelivery { + detached: StdArc>>, + } + + #[async_trait] + impl NativeDelivery for RecordingDelivery { + async fn deliver(&self, fact: NativeFact<'_>, _: HookCtx<'_>) -> Result<(), DeliveryError> { + if let NativeFact::Progress(NativeProgress::ToolDetached(result)) = fact { + self.detached.lock().unwrap().push(result.call_id.clone()); + } + Ok(()) + } + } + #[test] fn session_consumer_capabilities_are_typed_and_serde_defaulted() { let config = SessionConfig::new("session").with_response_attempt_supersession(); @@ -7165,13 +7749,15 @@ mod tests { let recorded = points.lock().unwrap().clone(); assert_eq!( - recorded.first(), - Some(&MutationPoint::AfterTurnEnded), - "first drive of a fresh turn must report AfterTurnEnded, got {recorded:?}" - ); - assert!( - recorded.contains(&MutationPoint::AfterToolResult), - "post-tool continuation must report AfterToolResult, got {recorded:?}" + recorded, + [ + MutationPoint::TurnStarted, + MutationPoint::AfterTurnEnded, + MutationPoint::AfterToolResult, + ], + "turn creation reports TurnStarted once, then the fresh turn's \ + first inference reports AfterTurnEnded and the post-tool \ + continuation reports AfterToolResult" ); } @@ -7974,6 +8560,7 @@ mod tests { release.clone(), "detached-done", )); + let delivery = RecordingDelivery::default(); let agent = Agent::builder() .model(FakeAdapter) .add_tool_source(tools) @@ -7982,6 +8569,7 @@ mod tests { .observer(RecordingObserver { events: events.clone(), }) + .delivery(delivery.clone()) .build() .unwrap(); @@ -8027,6 +8615,11 @@ mod tests { other => panic!("unexpected resumed step: {other:?}"), } + assert_eq!( + delivery.detached.lock().unwrap().as_slice(), + [ToolCallId::new("call-1")], + "the loop-authored detach placeholder is delivered as progress, once" + ); let events = events.lock().unwrap(); assert!(events.iter().any(|event| matches!( event, @@ -9070,6 +9663,7 @@ mod tests { )); driver .finish_cancelled(agentkit_core::TurnId::new("turn-cancelled"), Vec::new()) + .await .unwrap(); assert!( @@ -9285,7 +9879,7 @@ mod tests { } other => panic!("unexpected loop step: {other:?}"), }; - driver.cancel_pending_approval_for(call_id).unwrap(); + driver.cancel_pending_approval_for(call_id).await.unwrap(); assert!(driver.lifecycle.active_turn.is_none()); assert!(driver.pending_approvals.is_empty()); @@ -10250,6 +10844,10 @@ mod tests { fn on_transcript_event(&self, event: TranscriptEvent<'_>) { self.items.lock().unwrap().push(event.item.clone()); } + + fn on_transcript_rewrite(&self, event: TranscriptRewriteEvent<'_>) { + *self.items.lock().unwrap() = event.items.to_vec(); + } } #[tokio::test] diff --git a/docs/architecture.md b/docs/architecture.md index f55d433..e2d4890 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -138,7 +138,7 @@ stdio or a custom transport. ## Mutator path -Transcript edits — compaction, redaction, repair — plug into the loop through one generic seam, `LoopMutator`. Mutators run at well-defined `MutationPoint`s (`AfterToolResult`, `AfterTurnEnded`) and decide for themselves whether to touch the transcript via a `TranscriptCursor`. +Transcript edits — compaction, redaction, repair — plug into the loop through one generic seam, `LoopMutator`. Mutators run at well-defined `MutationPoint`s (`TurnStarted`, `AfterToolResult`, `AfterTurnEnded`) and decide for themselves whether to touch the transcript via a `TranscriptCursor`. For each registered mutator the loop: @@ -146,6 +146,6 @@ For each registered mutator the loop: 2. runs the mutator with read-only context and a mutable cursor over the transcript 3. emits `AgentEvent::MutationFinished { mutator, dirty, metadata, .. }` -If any mutator in the pass dirtied the cursor, the loop validates transcript invariants (tool_use ↔ tool_result pairing) and hard-fails with `LoopError::Mutator` if a mutator left the transcript protocol-invalid. +The pass is transactional: the chain edits a candidate copy, and if the final candidate differs from the live transcript the loop validates its invariants (tool_use ↔ tool_result pairing), re-checks cancellation, assigns it in one synchronous step, and delivers one `TranscriptObserver::on_transcript_rewrite` with the complete canonical transcript. A mutator that errors, left the transcript protocol-invalid, cancelled, or whose future was dropped changes nothing and notifies nobody; `LoopError::Mutator` surfaces the protocol violation. Compaction is the canonical mutator: `agentkit-compaction` provides a `Compactor` trait, `StrategyCompactor`, and trigger helpers (`item_count_trigger`, `context_window_trigger`) that wire into this seam through the `AgentBuilderCompactorExt::compactor` extension. Semantic compaction is provided through an injected `CompactionBackend` (e.g. `AgentCompactor`, which runs a nested sub-agent) rather than a built-in model client. diff --git a/docs/compaction.md b/docs/compaction.md index cb4cba5..667e43d 100644 --- a/docs/compaction.md +++ b/docs/compaction.md @@ -90,7 +90,9 @@ When a compactor fires, `agentkit-loop` emits: - `AgentEvent::MutationStarted { mutator, point, .. }` - `AgentEvent::MutationFinished { mutator, dirty, metadata, .. }` -Compactors run at every `MutationPoint` — currently `AfterToolResult` (between tool results and the next inference call) and `AfterTurnEnded`. The trigger closure decides which points are relevant. If a mutator's edit leaves the transcript protocol-invalid (orphaned/duplicate tool_use or tool_result), the loop hard-fails with `LoopError::Mutator` rather than letting the next request blow up at the provider. +Compactors run at every `MutationPoint` — currently `TurnStarted` (once per logical turn the driver creates, with queued input already appended), `AfterToolResult` (between tool results and the next inference call) and `AfterTurnEnded`. The trigger closure decides which points are relevant; one that ignores `point` now also fires at `TurnStarted`, compacting slightly earlier in the turn. If a mutator's edit leaves the transcript protocol-invalid (orphaned/duplicate tool_use or tool_result), the loop hard-fails with `LoopError::Mutator` rather than letting the next request blow up at the provider. + +The chain runs against a candidate copy of the transcript, so a failed, invalid, cancelled or dropped mutation leaves the live transcript unchanged. A committed change is published to every `TranscriptObserver` as one `on_transcript_rewrite` carrying the complete canonical transcript. ## Scope diff --git a/docs/loop.md b/docs/loop.md index 6fbf9b8..1633b8c 100644 --- a/docs/loop.md +++ b/docs/loop.md @@ -500,7 +500,7 @@ Recommended execution flow: 1. host submits input items 2. driver merges input into working transcript -3. driver runs registered `LoopMutator`s at `MutationPoint::AfterTurnEnded` / `AfterToolResult` (compaction, redaction, repair); the loop validates transcript invariants if any mutator dirtied the cursor +3. driver runs registered `LoopMutator`s at `MutationPoint::TurnStarted` / `AfterTurnEnded` / `AfterToolResult` (compaction, redaction, repair); the chain runs against a candidate copy, and a committed change is validated and published as one `TranscriptObserver::on_transcript_rewrite` 4. driver constructs `TurnRequest` 5. driver starts a provider turn 6. driver forwards streamed deltas to observers diff --git a/examples/openrouter-session-persistence/src/main.rs b/examples/openrouter-session-persistence/src/main.rs index b7bc834..1a780aa 100644 --- a/examples/openrouter-session-persistence/src/main.rs +++ b/examples/openrouter-session-persistence/src/main.rs @@ -5,7 +5,8 @@ //! - **Restore**: load prior `Item`s from sqlite and pass them to //! [`AgentBuilder::transcript`]. //! - **Incremental write**: register a [`TranscriptObserver`] that writes -//! each newly-appended item to sqlite as the loop runs. +//! each newly-appended item to sqlite as the loop runs, and replaces the +//! stored transcript when a mutator rewrites it. //! - **Resume**: the first `next()` call yields `AwaitingInput`, so a fresh //! process can pick up exactly where the previous one stopped — the //! transcript carries the prior turns, and the host supplies the next @@ -26,7 +27,7 @@ use std::sync::{Arc, Mutex}; use agentkit_core::{CancellationController, Item, ItemKind, Part}; use agentkit_loop::{ Agent, InputRequest, LoopInterrupt, LoopStep, PromptCacheRequest, PromptCacheRetention, - SessionConfig, TranscriptEvent, TranscriptObserver, + SessionConfig, TranscriptEvent, TranscriptObserver, TranscriptRewriteEvent, }; use agentkit_provider_openrouter::{OpenRouterAdapter, OpenRouterConfig}; use rusqlite::{Connection, params}; @@ -252,20 +253,42 @@ impl SqliteSessionStore { )?; Ok(()) } + + /// Replace the whole stored transcript with a canonical snapshot. + /// + /// A mutator (compaction, redaction, repair) rewrites history rather than + /// appending to it, so there is nothing to append — the stored rows have + /// to become the snapshot. + fn replace(&self, session_id: &str, items: &[Item]) -> Result<(), Box> { + let mut conn = self.conn.lock().unwrap(); + let tx = conn.transaction()?; + tx.execute( + "DELETE FROM items WHERE session_id = ?", + params![session_id], + )?; + for (seq, item) in items.iter().enumerate() { + tx.execute( + "INSERT INTO items (session_id, seq, json) VALUES (?, ?, ?)", + params![session_id, seq as i64, serde_json::to_string(item)?], + )?; + } + tx.commit()?; + Ok(()) + } } -/// `TranscriptObserver` impl that mirrors every appended item into sqlite. +/// `TranscriptObserver` impl that mirrors the driver's transcript into sqlite. /// -/// `on_transcript_event` is called synchronously by the loop and is the single -/// mutation point for the transcript -- every push funnels through here. The -/// observer must NOT block (sqlite writes are fast and local; for remote -/// stores, use a buffered channel and persist on a background task). +/// Both methods are called synchronously by the loop, and between them they +/// cover every way the transcript changes: `on_transcript_event` once per +/// appended item, `on_transcript_rewrite` with a complete snapshot when a +/// mutator rewrote history. Neither must block (sqlite writes are fast and +/// local; for remote stores, use a buffered channel and persist on a +/// background task). /// -/// Compaction-driven rewrites do **not** fire `on_transcript_event`. A -/// compaction-aware persistor would also subscribe to -/// `AgentEvent::CompactionFinished` via a `LoopObserver` and replace the -/// stored transcript when it sees that event. This example skips -/// compaction. +/// Both are best-effort post-commit observation, not a durability gate: the +/// loop has already committed by the time they run, so a failed write is +/// reported here and does not roll the turn back. struct SqliteTranscriptObserver { store: Arc, session_id: String, @@ -277,4 +300,10 @@ impl TranscriptObserver for SqliteTranscriptObserver { eprintln!("[persistence] failed to append item: {error}"); } } + + fn on_transcript_rewrite(&self, event: TranscriptRewriteEvent<'_>) { + if let Err(error) = self.store.replace(&self.session_id, event.items) { + eprintln!("[persistence] failed to replace transcript: {error}"); + } + } }