diff --git a/docs/user/agent-hooks-protocol.md b/docs/user/agent-hooks-protocol.md index 2a6e9fa8..5d01e5f4 100644 --- a/docs/user/agent-hooks-protocol.md +++ b/docs/user/agent-hooks-protocol.md @@ -38,6 +38,11 @@ global hook with the same `id`; other hooks from both files stay active. Kit skips an unreadable or malformed file with a warning. A subscription selector that matches no event Kit emits disables hooks with a warning. +Stdio hook processes inherit Kit's full environment, including any credentials +or other sensitive values present in the environment. +A transport's `env` entries override inherited values. Only configure stdio +hooks that you trust with those values. + Hook backends run local processes or receive session content. Review both files before you run Kit in a project you do not trust. @@ -46,6 +51,8 @@ before you run Kit in a project you do not trust. | Event | Where Kit emits it | Hook authority | | --- | --- | --- | | `session.start` | A session's model connection opens | Observe | +| `turn.start` | A logical agent-loop turn starts, with native session and turn IDs | Observe | +| `turn.end` | That logical turn completes, fails, or is cancelled | Observe | | `session.end` | A prompt run finishes or an ACP session closes | Observe | | `user.message.inbound` | A new turn admits a user prompt | Deny, replace the prompt | | `model.request.before` | Before each model call, with the transcript | Deny | @@ -70,5 +77,13 @@ Tool events cover the tools that compose scripts call, such as `shell`, `edit`, tool calls and tool results into message parts. It omits media, files and reasoning. -Kit does not emit `turn.*`, `model.switch.*`, `config.change.*`, `file.changed` +Tool calls, model retries and model continuations stay within the same logical +turn; they do not emit extra turn boundaries. Kit delivers turn observations +in order, before the session's next hook event and whenever the session goes +idle. An abrupt process exit can drop a terminal observation. The start event has no transcript snapshot; +its trigger is reported as external because the loop does not expose a cause. +The continuation count is zero because Kit does not initiate AHP hook +continuations. + +Kit does not emit `turn.finish.before`, `model.switch.*`, `config.change.*`, `file.changed` or the tool permission and progress events yet. diff --git a/mise.toml b/mise.toml index 07a0a911..b3d22974 100644 --- a/mise.toml +++ b/mise.toml @@ -58,11 +58,11 @@ run = [ [tasks.test] description = "Rust tests; arguments pass through, e.g. mise run test -- --test cli or mise run test -- --test runtime some_case" -run = "cargo test --locked" +run = "scripts/isolated-home.sh cargo test --locked" [tasks."test:headless"] description = "Rust tests without the tui feature; arguments pass through like test" -run = "cargo test --locked --no-default-features" +run = "scripts/isolated-home.sh cargo test --locked --no-default-features" [tasks."test:scripts"] description = "Shell tests for the release scripts" diff --git a/scripts/isolated-home.sh b/scripts/isolated-home.sh new file mode 100755 index 00000000..ce4dde20 --- /dev/null +++ b/scripts/isolated-home.sh @@ -0,0 +1,11 @@ +#!/bin/sh +# Runs a command with an empty HOME so tests never load the developer's +# global ~/.agents hooks, skills or other user configuration. +set -eu + +home=$(mktemp -d) +trap 'rm -rf "$home"' EXIT +CARGO_HOME=${CARGO_HOME:-$HOME/.cargo} \ + RUSTUP_HOME=${RUSTUP_HOME:-$HOME/.rustup} \ + HOME=$home \ + "$@" diff --git a/src/ahp.rs b/src/ahp.rs index 73c29055..7726749e 100644 --- a/src/ahp.rs +++ b/src/ahp.rs @@ -2,6 +2,7 @@ mod model; mod prompt; mod tool; +mod turn; use std::path::Path; use std::sync::Arc; @@ -28,13 +29,15 @@ use agentkit_loop::LoopError; pub(crate) use model::{HookedAdapter, HookedSession}; pub(crate) use prompt::{PromptHook, PromptRecord}; pub(crate) use tool::hooked_tools; +pub(crate) use turn::TurnHook; const SOURCE: &str = "urn:dev:kit:harness"; const CONFIG_FILE: &str = "hooks.json"; /// Registered hook backends with the event grants Kit implements. pub struct Hooks { - sdk: agenthooksprotocol::Hooks, + sdk: Arc, + turn_events: turn::TurnQueues, } impl Hooks { @@ -63,7 +66,10 @@ impl Hooks { fn new(registration: Registration, options: HooksOptions) -> Result { let sdk = agenthooksprotocol::Hooks::new(registration, options) .map_err(|error| error.to_string())?; - Ok(Self { sdk }) + Ok(Self { + sdk: Arc::new(sdk), + turn_events: turn::TurnQueues::default(), + }) } pub(crate) async fn session_start(&self, session: &str, trigger: SessionStartInputTrigger) { @@ -77,6 +83,8 @@ impl Hooks { } pub(crate) async fn session_end(&self, session: &str, outcome: SessionEndInputOutcome) { + self.drain_turn_events(session).await; + self.forget_turn_events(session); let reason = outcome.as_str().to_owned(); let input = SessionEndInput::new(outcome, reason, Box::new(Session::new(session))); settle(EventType::SessionEnd, self.sdk.session_end(input).await); @@ -164,6 +172,8 @@ fn options() -> Result { let capabilities = HostCapabilities { session_start: Some(capability::session_start().observe()), session_end: Some(capability::session_end().observe()), + turn_start: Some(capability::turn_start().observe()), + turn_end: Some(capability::turn_end().observe()), user_message_inbound: Some( capability::user_message_inbound() .deny() @@ -190,11 +200,13 @@ fn options() -> Result { ), ..HostCapabilities::default() }; - Ok(HooksOptions::new( + let mut options = HooksOptions::new( SOURCE, Harness::new("kit", env!("CARGO_PKG_VERSION")), capabilities, - )) + ); + options.backend.inherit_environment = true; + Ok(options) } fn read(path: &Path) -> Result, String> { diff --git a/src/ahp/model.rs b/src/ahp/model.rs index fc245fc3..f45c53c6 100644 --- a/src/ahp/model.rs +++ b/src/ahp/model.rs @@ -113,6 +113,9 @@ impl ModelSession for HookedSession { let inner = self.inner.begin_turn(request, cancellation).await?; return Ok(HookedTurn { inner, call: None }); }; + hooks + .drain_turn_events(&request.session_id.to_string()) + .await; self.attempts += 1; let Some(number) = Integer::new(self.attempts.into()) else { let inner = self.inner.begin_turn(request, cancellation).await?; diff --git a/src/ahp/prompt.rs b/src/ahp/prompt.rs index 6e41c1e7..7ac23a6f 100644 --- a/src/ahp/prompt.rs +++ b/src/ahp/prompt.rs @@ -62,6 +62,9 @@ impl LoopMutator for PromptHook { if ctx.point != MutationPoint::AfterTurnEnded { return Ok(()); } + self.hooks + .drain_turn_events(&ctx.session_id.to_string()) + .await; let Some(prompt) = cursor.last().filter(|item| item.kind == ItemKind::User) else { return Ok(()); }; diff --git a/src/ahp/test_support.rs b/src/ahp/test_support.rs index 36cb0474..6e308112 100644 --- a/src/ahp/test_support.rs +++ b/src/ahp/test_support.rs @@ -22,6 +22,10 @@ impl Backend { self.effects.lock().unwrap().insert(event.into(), effects); } + pub(crate) fn all_events(&self) -> Vec { + self.events.lock().unwrap().clone() + } + pub(crate) fn events(&self, event: &str) -> Vec { self.events .lock() @@ -74,6 +78,13 @@ pub(crate) fn observe(events: &[&str]) -> Value { pub(crate) fn hooks(subscriptions: Vec) -> (Arc, Arc) { let backend = Arc::new(Backend::default()); + (hooks_with_backend(subscriptions, backend.clone()), backend) +} + +pub(crate) fn hooks_with_backend( + subscriptions: Vec, + backend: Arc, +) -> Arc { let registration = serde_json::from_value(json!({ "protocolVersion": "draft", "hooks": [{ @@ -83,11 +94,8 @@ pub(crate) fn hooks(subscriptions: Vec) -> (Arc, Arc) { }], })) .unwrap(); - let options = options().unwrap().with_backend(BACKEND, backend.clone()); - ( - Arc::new(Hooks::new(registration, options).unwrap()), - backend, - ) + let options = options().unwrap().with_backend(BACKEND, backend); + Arc::new(Hooks::new(registration, options).unwrap()) } pub(crate) fn text_message(text: &str) -> Value { diff --git a/src/ahp/tests.rs b/src/ahp/tests.rs index f8cd3f2f..a42ebbd6 100644 --- a/src/ahp/tests.rs +++ b/src/ahp/tests.rs @@ -19,6 +19,63 @@ use serde_json::{Value, json}; use super::test_support::{hooks, intercept, observe, text_message}; use super::*; +#[cfg(unix)] +#[tokio::test] +async fn stdio_hooks_inherit_environment_with_explicit_overrides() { + const VARIABLE: &str = "KIT_AHP_TEST_INHERITED_ENV"; + const VALUE: &str = "environment-sentinel"; + // Configure only an isolated test process; other tests spawn hooks too. + if std::env::var(VARIABLE).ok().as_deref() != Some(VALUE) { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "ahp::tests::stdio_hooks_inherit_environment_with_explicit_overrides", + "--nocapture", + ]) + .env(VARIABLE, VALUE) + .env("HERDR_ENV", "1") + .env("HERDR_BIN_PATH", "herdr-bin-sentinel") + .env("HERDR_PANE_ID", "host-pane") + .env("HERDR_SOCKET_PATH", "socket-sentinel") + .output() + .unwrap(); + assert!(output.status.success(), "{output:?}"); + return; + } + + let directory = tempfile::tempdir().unwrap(); + let output = directory.path().join("hook-env"); + let registration = serde_json::from_value(json!({ + "protocolVersion": "draft", + "hooks": [{ + "id": "dev.kit.environment-test", + "transport": { + "type": "stdio", + "command": "/bin/sh", + "args": [ + "-c", + "read -r event; printf '%s\\n' \"$HERDR_ENV\" \"$HERDR_BIN_PATH\" \"$HERDR_PANE_ID\" \"$HERDR_SOCKET_PATH\" \"$KIT_AHP_TEST_INHERITED_ENV\" > \"$1\"", + "environment-hook", + output, + ], + "env": {"HERDR_PANE_ID": "configured-pane"}, + "lifecycle": "per_event", + }, + "subscriptions": [observe(&["session.start"])], + }], + })) + .unwrap(); + let hooks = Hooks::new(registration, options().unwrap()).unwrap(); + hooks + .session_start("environment-test", SessionStartInputTrigger::Startup) + .await; + hooks.sdk.shutdown().await.unwrap(); + assert_eq!( + std::fs::read_to_string(output).unwrap(), + "1\nherdr-bin-sentinel\nconfigured-pane\nsocket-sentinel\nenvironment-sentinel\n", + ); +} + #[test] fn every_granted_event_accepts_wildcard_subscriptions() { hooks(vec![intercept(&["*"]), observe(&["*"])]); @@ -56,6 +113,8 @@ async fn every_lifecycle_event_reports_harness_identity() { for event in [ "session.start", "session.end", + "turn.start", + "turn.end", "user.message.inbound", "model.request.before", "model.response.after", @@ -475,6 +534,7 @@ async fn admit(hooks: &Arc, model: &Model, input: Vec) -> (LoopStep let opened = crate::session::open_in(root.path(), &directory, "abc", false, false, initial).unwrap(); let mut driver = Agent::builder() + .observer(TurnHook(hooks.clone())) .model(HookedAdapter::new( model.clone(), Some(hooks.clone()), @@ -494,6 +554,7 @@ async fn admit(hooks: &Arc, model: &Model, input: Vec) -> (LoopStep .await .unwrap(); let step = driver.next().await.unwrap(); + hooks.drain_turn_events("abc").await; let persisted = crate::session::load_in(root.path(), &directory, "abc").unwrap(); (step, persisted) } @@ -650,3 +711,277 @@ async fn compaction_hooks_supply_and_rewrite_the_summary() { "rewritten" ); } + +#[tokio::test] +async fn turn_boundaries_keep_native_ids_across_turns_and_sessions() { + let (hooks, backend) = hooks(vec![observe(&["turn.start", "turn.end"])]); + for session in ["first", "second"] { + let mut driver = Agent::builder() + .model(Model::default()) + .observer(TurnHook(hooks.clone())) + .input(vec![Item::text(ItemKind::User, "one")]) + .build() + .unwrap() + .start(SessionConfig::new(session)) + .await + .unwrap(); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Finished(_) + )); + driver + .submit_input(vec![Item::text(ItemKind::User, "two")]) + .unwrap(); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Finished(_) + )); + hooks.drain_turn_events(session).await; + } + let starts = backend.events("turn.start"); + let ends = backend.events("turn.end"); + assert_eq!(starts.len(), 4); + assert_eq!(ends.len(), 4); + for (start, end) in starts.iter().zip(&ends) { + assert_eq!(start["session"], end["session"]); + assert_eq!(start["turn"], end["turn"]); + assert_eq!(end["outcome"], "completed"); + } + assert_eq!(starts[0]["session"]["id"], "first"); + assert_eq!(starts[2]["session"]["id"], "second"); + assert_ne!(starts[0]["turn"]["id"], starts[1]["turn"]["id"]); +} + +/// Holds `turn.end` at the hook transport until the test releases it. +struct HoldTurnEnd { + events: Arc, + reached: tokio::sync::Notify, + release: tokio::sync::Notify, +} + +impl agenthooksprotocol::adapters::registered::ManagedBackend for HoldTurnEnd { + fn call( + &self, + request: Value, + timeout: std::time::Duration, + ) -> agenthooksprotocol::client::LocalFuture< + '_, + Result, + > { + Box::pin(async move { + if request["params"]["event"]["type"] == "turn.end" { + self.reached.notify_one(); + self.release.notified().await; + } + self.events.call(request, timeout).await + }) + } + + fn shutdown( + &self, + ) -> agenthooksprotocol::client::LocalFuture< + '_, + Result<(), agenthooksprotocol::client::HookError>, + > { + self.events.shutdown() + } +} + +#[tokio::test] +async fn cancelled_drain_still_delivers_turn_end_before_session_end() { + let backend = Arc::new(HoldTurnEnd { + events: Arc::default(), + reached: tokio::sync::Notify::new(), + release: tokio::sync::Notify::new(), + }); + let hooks = super::test_support::hooks_with_backend( + vec![observe(&["turn.start", "turn.end", "session.end"])], + backend.clone(), + ); + let mut driver = Agent::builder() + .model(Model::default()) + .observer(TurnHook(hooks.clone())) + .input(vec![Item::text(ItemKind::User, "hello")]) + .build() + .unwrap() + .start(SessionConfig::new("held")) + .await + .unwrap(); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Finished(_) + )); + tokio::select! { + () = hooks.drain_turn_events("held") => panic!("turn.end was not held"), + () = backend.reached.notified() => {} + } + backend.release.notify_one(); + hooks + .session_end("held", SessionEndInputOutcome::Completed) + .await; + let types: Vec<_> = backend + .events + .all_events() + .iter() + .map(|event| event["type"].as_str().unwrap().to_owned()) + .collect(); + assert_eq!(types, ["turn.start", "turn.end", "session.end"]); +} + +#[tokio::test] +async fn failed_model_request_ends_the_logical_turn() { + let (hooks, backend) = hooks(vec![ + observe(&["turn.start", "turn.end"]), + intercept(&["model.request.before"]), + ]); + backend.respond( + "model.request.before", + json!([{ "type": "deny", "reason": "blocked" }]), + ); + let mut driver = Agent::builder() + .model(HookedAdapter::new( + Model::default(), + Some(hooks.clone()), + SessionStartInputTrigger::Startup, + )) + .observer(TurnHook(hooks.clone())) + .input(vec![Item::text(ItemKind::User, "hello")]) + .build() + .unwrap() + .start(SessionConfig::new("failed")) + .await + .unwrap(); + assert!(driver.next().await.is_err()); + hooks.drain_turn_events("failed").await; + let starts = backend.events("turn.start"); + let ends = backend.events("turn.end"); + assert_eq!(starts.len(), 1); + assert_eq!(ends.len(), 1); + assert_eq!(ends[0]["outcome"], "failed"); + assert_eq!(starts[0]["turn"], ends[0]["turn"]); + assert_eq!(ends[0]["session"]["id"], "failed"); +} + +#[derive(Clone)] +struct ToolRoundModel(bool); + +#[async_trait] +impl ModelAdapter for ToolRoundModel { + type Session = Self; + async fn start_session(&self, _: SessionConfig) -> Result { + Ok(self.clone()) + } +} + +#[async_trait] +impl ModelSession for ToolRoundModel { + type Turn = ToolRoundReply; + async fn begin_turn( + &mut self, + _: TurnRequest, + _: Option, + ) -> Result { + let mut events = std::collections::VecDeque::new(); + let call = agentkit_core::ToolCallPart::new("echo-call", "echo", json!({"text": "hello"})); + let mut assistant = Item::text(ItemKind::Assistant, ""); + assistant.parts = vec![Part::ToolCall(call.clone())]; + if self.0 { + // A superseded provider response is retried within this model turn. + events.push_back(ModelTurnEvent::ToolCall(agentkit_core::ToolCallPart::new( + "discarded-call", + "echo", + json!({"text": "discarded"}), + ))); + events.push_back(ModelTurnEvent::ResponseAttemptSuperseded); + events.push_back(ModelTurnEvent::ToolCall(call)); + } + events.push_back(ModelTurnEvent::Finished(ModelTurnResult { + finish_reason: if self.0 { + FinishReason::ToolCall + } else { + FinishReason::Completed + }, + output_items: if self.0 { + vec![assistant] + } else { + vec![Item::text(ItemKind::Assistant, "done")] + }, + usage: None, + metadata: MetadataMap::new(), + model: None, + response_id: None, + })); + self.0 = false; + Ok(ToolRoundReply(events)) + } +} + +struct ToolRoundReply(std::collections::VecDeque); +#[async_trait] +impl ModelTurn for ToolRoundReply { + async fn next_event( + &mut self, + _: Option, + ) -> Result, LoopError> { + Ok(self.0.pop_front()) + } +} + +#[tokio::test] +async fn retry_and_tool_continuation_are_one_turn_and_retirement_cancels_once() { + for retire in [false, true] { + let (hooks, backend) = hooks(vec![ + observe(&["turn.start", "turn.end"]), + intercept(&["model.request.before"]), + ]); + let echo = Echo::new(); + let mut registry = ToolRegistry::new(); + registry.register_arc(echo.clone()); + let mut driver = Agent::builder() + .model(HookedAdapter::new( + ToolRoundModel(true), + Some(hooks.clone()), + SessionStartInputTrigger::Startup, + )) + .observer(TurnHook(hooks.clone())) + .add_tool_source(registry) + .input(vec![Item::text(ItemKind::User, "hello")]) + .build() + .unwrap() + .start(SessionConfig::new("tools")) + .await + .unwrap(); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Interrupt(agentkit_loop::LoopInterrupt::AfterToolResult(_)) + )); + hooks.drain_turn_events("tools").await; + assert_eq!(backend.events("turn.start").len(), 1); + assert!(backend.events("turn.end").is_empty()); + assert_eq!(echo.received.lock().unwrap().len(), 1); + if retire { + assert!(driver.retire_interrupted_turn().await.unwrap().is_some()); + assert!(driver.retire_interrupted_turn().await.unwrap().is_none()); + } else { + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Finished(_) + )); + } + hooks.drain_turn_events("tools").await; + let starts = backend.events("turn.start"); + let ends = backend.events("turn.end"); + assert_eq!(starts.len(), 1); + assert_eq!(ends.len(), 1); + assert_eq!(starts[0]["turn"], ends[0]["turn"]); + assert_eq!(ends[0]["session"]["id"], "tools"); + assert_eq!( + ends[0]["outcome"], + if retire { "cancelled" } else { "completed" } + ); + assert_eq!( + backend.events("model.request.before").len(), + if retire { 1 } else { 2 } + ); + } +} diff --git a/src/ahp/turn.rs b/src/ahp/turn.rs new file mode 100644 index 00000000..c3b3aa06 --- /dev/null +++ b/src/ahp/turn.rs @@ -0,0 +1,119 @@ +use std::collections::{HashMap, VecDeque}; +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; + +use agenthooksprotocol::inputs::{TurnEndInput, TurnStartInput}; +use agenthooksprotocol::{ + EventType, Integer, Session, TurnEndInputOutcome, TurnEndInputTurn, TurnStartInputTrigger, + TurnStartInputTurn, +}; +use agentkit_core::FinishReason; +use agentkit_loop::{AgentEvent, LoopObserver, ObservedEvent}; + +use super::Hooks; + +pub(super) type TurnQueues = Arc>>; + +/// Queues native turn boundaries until the session's next async hook point. +#[derive(Clone)] +pub(crate) struct TurnHook(pub(crate) Arc); + +#[derive(Default)] +pub(super) struct SessionTurns { + pending: VecDeque, + delivery: Arc>, +} + +impl Hooks { + /// Delivers the turn observations queued for `session`, in order. + pub(crate) async fn drain_turn_events(&self, session: &str) { + let delivery = Arc::clone( + &lock(&self.turn_events) + .entry(session.to_owned()) + .or_default() + .delivery, + ); + let guard = delivery.lock_owned().await; + if lock(&self.turn_events) + .get(session) + .is_none_or(|turns| turns.pending.is_empty()) + { + return; + } + let sdk = Arc::clone(&self.sdk); + let queues = Arc::clone(&self.turn_events); + let session = session.to_owned(); + // A cancelled caller must not abandon an observation the SDK has begun delivering. + let delivered = tokio::spawn(async move { + let _guard = guard; + while let Some(event) = next(&queues, &session) { + deliver(&sdk, &session, event).await; + } + }); + if delivered.await.is_err() { + tracing::warn!("agent hook turn observation delivery failed"); + } + } + + pub(super) fn forget_turn_events(&self, session: &str) { + lock(&self.turn_events).remove(session); + } +} + +fn lock(queues: &TurnQueues) -> MutexGuard<'_, HashMap> { + queues.lock().unwrap_or_else(PoisonError::into_inner) +} + +fn next(queues: &TurnQueues, session: &str) -> Option { + lock(queues).get_mut(session)?.pending.pop_front() +} + +async fn deliver(sdk: &agenthooksprotocol::Hooks, session: &str, event: AgentEvent) { + let owner = Box::new(Session::new(session)); + match event { + AgentEvent::TurnStarted { turn_id, .. } => { + let input = TurnStartInput::new( + Vec::new(), + TurnStartInputTrigger::External, + TurnStartInputTurn::new(turn_id.to_string()), + ) + .with_session(owner); + super::settle(EventType::TurnStart, sdk.turn_start(input).await); + } + AgentEvent::TurnFinished(result) => { + let outcome = match result.finish_reason { + FinishReason::Cancelled => TurnEndInputOutcome::Cancelled, + FinishReason::Error | FinishReason::Blocked => TurnEndInputOutcome::Failed, + _ => TurnEndInputOutcome::Completed, + }; + let Some(continuations) = Integer::from_i64(0) else { + return; + }; + let input = TurnEndInput::new( + continuations, + Vec::new(), + outcome, + TurnEndInputTurn::new(result.turn_id.to_string()), + ) + .with_session(owner) + .with_sources() + .with_items(super::messages(&result.items)); + super::settle(EventType::TurnEnd, sdk.turn_end(input).await); + } + _ => {} + } +} + +impl LoopObserver for TurnHook { + fn handle_event(&self, observed: ObservedEvent) { + if matches!( + observed.event, + AgentEvent::TurnStarted { .. } | AgentEvent::TurnFinished(_) + ) { + lock(&self.0.turn_events) + .entry(observed.session_id.to_string()) + .or_default() + .pending + .push_back(observed.event); + } + } +} diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index f1465e73..975fa2cb 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -2148,6 +2148,9 @@ async fn session_actor(actor: SessionActor) { None => break, }, } + if let Some(hooks) = runtime.hooks() { + hooks.drain_turn_events(&session_id.to_string()).await; + } } } diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 980adb69..7fd2c3ef 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -1551,6 +1551,9 @@ async fn session_actor(actor: SessionActor) None => break, }, } + if let Some(hooks) = runtime.hooks() { + hooks.drain_turn_events(&session_id.to_string()).await; + } } } diff --git a/src/runtime.rs b/src/runtime.rs index c56c47e7..b3fcd597 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -1426,6 +1426,7 @@ impl Runtime { .task_manager(task_manager) .mutator(compactor); if let Some(hooks) = &self.hooks { + builder = builder.observer(crate::ahp::TurnHook(Arc::clone(hooks))); builder = builder.mutator(crate::ahp::PromptHook::new( Arc::clone(hooks), "cli", @@ -1482,13 +1483,25 @@ impl Runtime { } let result = { let cancellation = std::pin::pin!(cancelled.cancelled()); - let run = std::pin::pin!(drive_with_tasks(&mut driver, Some(&tasks))); + let run = std::pin::pin!(drive_with_tasks( + &mut driver, + Some(&tasks), + self.hooks().map(|hooks| (hooks, session_id.as_str())), + )); // Idle prompts must observe cancellation too, not just model turns. match select(cancellation, run).await { Either::Left(((), _)) => Err(LoopError::Cancelled), Either::Right((result, _)) => result, } }; + // The outer cancellation select can drop the drive future before the + // loop emits its terminal event. Retire after that borrow is released; + // this is a no-op if the driver already finished the cancelled turn. + if matches!(&result, Err(LoopError::Cancelled)) + && let Err(error) = driver.retire_interrupted_turn().await + { + tracing::warn!(%error, "failed to retire cancelled prompt turn"); + } // Cancel owned compose work cooperatively before aborting task handles. background_jobs.cancel_all(); let _ = tokio::time::timeout( @@ -1592,11 +1605,13 @@ impl Runtime { .task_manager(background_task_manager()) .mutator(compactor); let builder = match &self.hooks { - Some(hooks) => builder.mutator(crate::ahp::PromptHook::new( - Arc::clone(hooks), - "a2a", - crate::ahp::PromptRecord::Transient(recorded.clone()), - )), + Some(hooks) => builder + .observer(crate::ahp::TurnHook(Arc::clone(hooks))) + .mutator(crate::ahp::PromptHook::new( + Arc::clone(hooks), + "a2a", + crate::ahp::PromptRecord::Transient(recorded.clone()), + )), None => builder, }; let builder = builder @@ -1818,6 +1833,7 @@ impl Runtime { .task_manager(task_manager) .mutator(compactor); if let Some(hooks) = &self.hooks { + builder = builder.observer(crate::ahp::TurnHook(Arc::clone(hooks))); builder = builder.mutator(crate::ahp::PromptHook::new( Arc::clone(hooks), "acp", @@ -2949,12 +2965,13 @@ fn record_loop_failure( async fn drive( driver: &mut LoopDriver>, ) -> Result { - drive_with_tasks(driver, None).await + drive_with_tasks(driver, None, None).await } async fn drive_with_tasks( driver: &mut LoopDriver, tasks: Option<&TaskManagerHandle>, + turn_events: Option<(&crate::ahp::Hooks, &str)>, ) -> Result { let mut output = String::new(); loop { @@ -2988,6 +3005,9 @@ async fn drive_with_tasks( if !running { return Ok(output); } + if let Some((hooks, session)) = turn_events { + hooks.drain_turn_events(session).await; + } if let Some(tasks) = tasks { // Events are queued during model turns, including failure // and cancellation; no check-to-wait notification is lost. diff --git a/src/runtime/prompt_tests.rs b/src/runtime/prompt_tests.rs index 5977c43c..b4d4e2b4 100644 --- a/src/runtime/prompt_tests.rs +++ b/src/runtime/prompt_tests.rs @@ -179,7 +179,7 @@ async fn last_completion_queued_during_final_turn_is_not_lost() { .unwrap(); let output = tokio::time::timeout( Duration::from_secs(3), - drive_with_tasks(&mut driver, Some(&tasks)), + drive_with_tasks(&mut driver, Some(&tasks), None), ) .await .unwrap() @@ -187,3 +187,116 @@ async fn last_completion_queued_during_final_turn_is_not_lost() { assert_eq!(output, "resumed final answer"); assert!(tasks.list_running().await.is_empty()); } + +/// Cancels at the real hook transport boundary while model admission is pending. +struct CancelModelAdmission { + events: Arc, + cancellation: CancellationController, +} + +impl agenthooksprotocol::adapters::registered::ManagedBackend for CancelModelAdmission { + fn call( + &self, + request: serde_json::Value, + timeout: Duration, + ) -> agenthooksprotocol::client::LocalFuture< + '_, + Result, + > { + Box::pin(async move { + let cancels = request["params"]["event"]["type"] == "model.request.before"; + let response = self.events.call(request, timeout).await?; + if cancels { + self.cancellation.interrupt(); + // Keep drive_with_tasks suspended until the outer select wins. + std::future::pending().await + } else { + Ok(response) + } + }) + } + + fn shutdown( + &self, + ) -> agenthooksprotocol::client::LocalFuture< + '_, + Result<(), agenthooksprotocol::client::HookError>, + > { + self.events.shutdown() + } +} + +#[tokio::test] +async fn persistent_cancellation_emits_cancelled_turn_before_session_end() { + const ISOLATED: &str = "KIT_AHP_PERSISTENT_CANCEL_TEST"; + if std::env::var(ISOLATED).ok().as_deref() != Some("1") { + let output = tokio::task::spawn_blocking(|| { + std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "runtime::prompt_tests::persistent_cancellation_emits_cancelled_turn_before_session_end", "--nocapture"]) + .env(ISOLATED, "1") + .env("OPENROUTER_API_KEY", "test-only-key") + // This endpoint has no models-catalog URL and must never be called. + .env("OPENROUTER_BASE_URL", "http://127.0.0.1:1/completions") + .output().unwrap() + }).await.unwrap(); + assert!(output.status.success(), "{output:?}"); + return; + } + + use crate::ahp::test_support::{Backend, hooks_with_backend, intercept, observe}; + let root = tempfile::tempdir().unwrap(); + let cancellation = CancellationController::new(); + let events = Arc::new(Backend::default()); + let hooks = hooks_with_backend( + vec![ + observe(&["turn.start", "turn.end", "session.end"]), + intercept(&["model.request.before"]), + ], + Arc::new(CancelModelAdmission { + events: events.clone(), + cancellation: cancellation.clone(), + }), + ); + let mut runtime = Runtime::with_session_provider_and_credentials( + root.path(), + "test/model", + ProviderKind::OpenRouter, + SessionRequest { + id: "persistent-cancellation".into(), + resume: false, + force: false, + }, + crate::credentials::CredentialStorage::Memory, + ) + .unwrap(); + Arc::get_mut(&mut runtime).unwrap().hooks = Some(hooks); + let result = tokio::time::timeout( + Duration::from_secs(10), + runtime.run_persistent_interruptible("hello".into(), Some(cancellation.handle())), + ) + .await + .unwrap(); + let error = result.unwrap_err(); + assert!(error.contains("cancel"), "{error}"); + + let delivered = events.all_events(); + let terminal: Vec<_> = delivered + .iter() + .filter(|event| { + matches!( + event["type"].as_str(), + Some("turn.start" | "turn.end" | "session.end") + ) + }) + .collect(); + assert_eq!(terminal.len(), 3, "{delivered:?}"); + assert_eq!(terminal[0]["type"], "turn.start"); + assert_eq!(terminal[1]["type"], "turn.end"); + assert_eq!(terminal[2]["type"], "session.end"); + assert_eq!(terminal[1]["outcome"], "cancelled"); + assert_eq!(terminal[2]["outcome"], "cancelled"); + assert_eq!(terminal[0]["turn"], terminal[1]["turn"]); + assert_eq!(terminal[0]["session"], terminal[1]["session"]); + assert_eq!(terminal[1]["session"], terminal[2]["session"]); + assert_eq!(events.events("model.request.before").len(), 1); +}