diff --git a/Cargo.lock b/Cargo.lock index b87131f3..ad69c8fb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -87,6 +87,22 @@ dependencies = [ "syn 3.0.4", ] +[[package]] +name = "agenthooksprotocol" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "036c046bd6b6f4b4311f7ecdcb57972eb44e738cf65196ebc8501b4fdfb135c8" +dependencies = [ + "jsonschema 0.33.0", + "reqwest 0.12.28", + "serde", + "serde_json", + "sha2 0.10.9", + "time", + "tokio", + "url", +] + [[package]] name = "agentkit-acp" version = "0.10.12" @@ -116,7 +132,7 @@ dependencies = [ "axum", "eventsource-stream", "futures", - "reqwest", + "reqwest 0.13.5", "serde_json", "thiserror 2.0.18", "tokio", @@ -178,7 +194,7 @@ dependencies = [ "async-trait", "base64 0.22.1", "futures-util", - "reqwest", + "reqwest 0.13.5", "serde", "serde_json", "thiserror 2.0.18", @@ -248,7 +264,7 @@ dependencies = [ "futures-util", "http", "httpdate", - "reqwest", + "reqwest 0.13.5", "serde", "serde_json", "thiserror 2.0.18", @@ -285,7 +301,7 @@ dependencies = [ "async-trait", "futures-util", "http", - "reqwest", + "reqwest 0.13.5", "rmcp", "serde", "serde_json", @@ -322,7 +338,7 @@ dependencies = [ "async-trait", "base64 0.22.1", "futures-util", - "reqwest", + "reqwest 0.13.5", "serde", "serde_json", "thiserror 2.0.18", @@ -2075,6 +2091,17 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "fluent-uri" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1918b65d96df47d3591bed19c5cca17e3fa5d0707318e4b5ef2eae01764df7e5" +dependencies = [ + "borrow-or-share", + "ref-cast", + "serde", +] + [[package]] name = "fluent-uri" version = "0.4.1" @@ -2107,6 +2134,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fraction" +version = "0.15.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e076045bb43dac435333ed5f04caf35c7463631d0dae2deb2638d94dd0a5b872" +dependencies = [ + "lazy_static", + "num", +] + [[package]] name = "fraction" version = "0.17.0" @@ -2972,6 +3009,32 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonschema" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d46662859bc5f60a145b75f4632fbadc84e829e45df6c5de74cfc8e05acb96b5" +dependencies = [ + "ahash", + "base64 0.22.1", + "bytecount", + "email_address", + "fancy-regex 0.16.2", + "fraction 0.15.4", + "idna", + "itoa", + "num-cmp", + "num-traits", + "once_cell", + "percent-encoding", + "referencing 0.33.0", + "regex", + "regex-syntax", + "serde", + "serde_json", + "uuid-simd", +] + [[package]] name = "jsonschema" version = "0.58.5" @@ -2983,7 +3046,7 @@ dependencies = [ "data-encoding", "email_address", "fancy-regex 0.19.0", - "fraction", + "fraction 0.17.0", "getrandom 0.3.4", "idna", "itoa", @@ -2992,7 +3055,7 @@ dependencies = [ "num-cmp", "num-traits", "percent-encoding", - "referencing", + "referencing 0.58.5", "regex", "serde", "serde_json", @@ -3018,7 +3081,7 @@ checksum = "4e3bfac11ac357ec620880025757d38c67da8fd5b9ef5024d19f0f360002b4ca" dependencies = [ "ahash", "bytecount", - "fraction", + "fraction 0.17.0", "getrandom 0.3.4", "num-cmp", "num-traits", @@ -3083,6 +3146,7 @@ dependencies = [ "a2a-protocol-client", "a2a-protocol-server", "a2a-protocol-types", + "agenthooksprotocol", "agentkit-acp", "agentkit-acp-http", "agentkit-acp-schema", @@ -3120,7 +3184,7 @@ dependencies = [ "hyper-util", "ignore", "image", - "jsonschema", + "jsonschema 0.58.5", "jsonwebtoken", "keyring", "libc", @@ -3131,7 +3195,7 @@ dependencies = [ "opus-pure", "ratatui", "ratatui-image", - "reqwest", + "reqwest 0.13.5", "rmcp", "rpassword", "runlet", @@ -3725,7 +3789,7 @@ dependencies = [ "bytes", "http", "opentelemetry", - "reqwest", + "reqwest 0.13.5", ] [[package]] @@ -3740,7 +3804,7 @@ dependencies = [ "opentelemetry-proto", "opentelemetry_sdk", "prost", - "reqwest", + "reqwest 0.13.5", "serde_json", "thiserror 2.0.18", "tokio", @@ -4443,6 +4507,20 @@ dependencies = [ "syn 3.0.4", ] +[[package]] +name = "referencing" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e9c261f7ce75418b3beadfb3f0eb1299fe8eb9640deba45ffa2cb783098697d" +dependencies = [ + "ahash", + "fluent-uri 0.3.2", + "once_cell", + "parking_lot", + "percent-encoding", + "serde_json", +] + [[package]] name = "referencing" version = "0.58.5" @@ -4450,7 +4528,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "590efadb0a669f1712c1e3ab810a55d4c3d20d0127c455f7236dd0710aba0346" dependencies = [ "ahash", - "fluent-uri", + "fluent-uri 0.4.1", "getrandom 0.3.4", "hashbrown 0.17.1", "itoa", @@ -4489,6 +4567,44 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tower", + "tower-http 0.6.11", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "webpki-roots", +] + [[package]] name = "reqwest" version = "0.13.5" @@ -4561,7 +4677,7 @@ dependencies = [ "oauth2", "pin-project-lite", "process-wrap", - "reqwest", + "reqwest 0.13.5", "serde", "serde_json", "sse-stream", @@ -6042,6 +6158,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "23b082222b4f6619906941c17eb2297fff4c2fb96cb60164170522942a200bd8" dependencies = [ "outref", + "uuid", "vsimd", ] diff --git a/Cargo.toml b/Cargo.toml index 19e5677d..4bf55564 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,7 @@ tokio-tungstenite = { version = "=0.30.0", default-features = false, features = a2a-protocol-client = "=0.11.0" a2a-protocol-server = { version = "=0.11.0", default-features = false } a2a-protocol-types = "=0.11.0" +agenthooksprotocol = { version = "=0.4.1", default-features = false, features = ["reqwest", "tokio-process"] } agent-client-protocol = { package = "agentkit-acp-sdk", version = "=2.2.0", features = ["unstable_session_fork", "unstable_session_inject"] } agent-client-protocol-schema = { package = "agentkit-acp-schema", version = "=1.9.1", features = ["unstable_session_notices", "unstable_session_compaction"] } agent-client-protocol-http = { package = "agentkit-acp-http", version = "=2.2.0", default-features = false, features = ["server", "client"] } diff --git a/docs/user/agent-hooks-protocol.md b/docs/user/agent-hooks-protocol.md new file mode 100644 index 00000000..2a6e9fa8 --- /dev/null +++ b/docs/user/agent-hooks-protocol.md @@ -0,0 +1,74 @@ +# Agent Hooks Protocol + +Kit delivers [Agent Hooks Protocol](https://github.com/agenthooksprotocol/rust-sdk) +events to hook backends registered in `~/.agents/hooks.json` (global) and +`/.agents/hooks.json` (project). When neither file exists, hooks are off. + +## Configure hooks + +Each file holds an AHP registration: + +```json +{ + "protocolVersion": "draft", + "hooks": [ + { + "id": "org.example.policy", + "transport": { + "type": "stdio", + "command": "/usr/local/bin/policy-hook", + "lifecycle": "persistent" + }, + "subscriptions": [ + { + "events": ["tool.before"], + "mode": "intercept", + "timeoutMs": 3000, + "failurePolicy": "fail-closed", + "content": {"default": "body"} + } + ] + } + ] +} +``` + +Kit reads the global file, then the project file. A project hook replaces the +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. + +Hook backends run local processes or receive session content. Review both files +before you run Kit in a project you do not trust. + +## Events + +| Event | Where Kit emits it | Hook authority | +| --- | --- | --- | +| `session.start` | A session's model connection opens | 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 | +| `model.response.after` | A model call finishes, with its output | Observe | +| `model.error` | A model call fails | Observe | +| `tool.before` | Before a tool runs | Allow, deny, return a result, replace or merge input | +| `tool.after` | A tool finishes, with its output | Observe | +| `context.compact.before` | Before compaction summarizes the transcript | Return a summary, replace summarizer instructions | +| `context.compact.after` | Before the summary is installed | Replace the summary | + +A denied prompt is removed, together with resources attached to it, and the +turn ends without a model call. A denied +model request fails the turn. A denied tool call returns a permission error to +the model. + +Kit checks edited tool input against the tool's input schema before the tool +runs; an edit that does not match counts as a hook failure. An edit that the +hook does not also allow goes through Kit's permission policy again. + +Tool events cover the tools that compose scripts call, such as `shell`, `edit`, +`read_file` and the MCP `tool` dispatcher. Kit projects text, structured data, +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` +or the tool permission and progress events yet. diff --git a/docs/user/getting-started-and-configuration.md b/docs/user/getting-started-and-configuration.md index c2e9adb6..2e6a91b8 100644 --- a/docs/user/getting-started-and-configuration.md +++ b/docs/user/getting-started-and-configuration.md @@ -288,6 +288,8 @@ subdir = "agent-plugins/example" `a2a` applies to `serve` and `tui`. Configured plugins can provide MCP servers without `mcp_config`; supported `stdio` and `streamable-http` declarations are registered, while `sse` declarations are skipped with a stderr diagnostic. MCP servers merge by name in this order: plugins, configured `mcp_config`, `/.mcp.json`, then `--mcp-config`. Higher layers replace whole conflicts while preserving non-conflicting lower entries; live removal reveals the next lower layer. Plugin data is stored under `/plugin-data/`. See [Agent Plugins](agent-plugins.md) for placeholders, collision rules, and ACP child behavior. `otel_endpoint` enables OTLP export of AgentKit's GenAI trace spans. `otel_protocol` accepts exactly `grpc`, `http/protobuf`, or `http/json`; gRPC is the default. For either HTTP protocol, Kit appends `/v1/traces` to a base endpoint unless it already ends with that path. `credential_store` selects one backend for OpenAI, Speakeasy, and MCP and defaults to `memory`; selecting `file` requires `credential_dir`, while a credential directory is invalid with `memory` or `keychain`. Memory credentials are process-local and are not shared with the TUI server process or nested Kit children. Standalone OpenAI and Speakeasy login requires persistent `keychain` or `file` storage. ACP profiles are direct executable-and-argument configurations, not shell command strings. `[subagent].harness` must name an available fully qualified profile such as `acp.review`; otherwise startup reports `unknown subagent ACP harness`. When no subagent harness is selected, the built-in `acp.kit` profile is used. +Agent hooks are configured in `~/.agents/hooks.json` and `/.agents/hooks.json`. See [Agent Hooks Protocol](agent-hooks-protocol.md). + ### Configuration precedence and built-in defaults For settings exposed by a command, precedence is: diff --git a/src/ahp.rs b/src/ahp.rs new file mode 100644 index 00000000..76af1897 --- /dev/null +++ b/src/ahp.rs @@ -0,0 +1,407 @@ +//! Agent Hooks Protocol integration at Kit's native lifecycle seams. +mod model; +mod prompt; +mod tool; + +use std::path::Path; +use std::sync::Arc; + +use agenthooksprotocol::capability::{self, ModifyOperation}; +use agenthooksprotocol::client::{ + ConsumerProtocolOutcome, DeliveryDiagnostic, HookError, SettledCandidate, +}; +use agenthooksprotocol::generated::{Effect, Event, Presence}; +use agenthooksprotocol::hooks::{EventOutcome, HooksOptions}; +use agenthooksprotocol::inputs::{ + ContextCompactAfterInput, ContextCompactBeforeInput, MessageInput, PartInput, SessionEndInput, + SessionStartInput, +}; +use agenthooksprotocol::{ + Attachment, CanonicalMessageRole, ContentItem, ContextCompactAfterInputRemovedItem, + ContextCompactBeforeInputTrigger, EventType, ExecutionEventExecution, + ExecutionEventExecutionExecuted, Registration, Session, SessionEndInputOutcome, + SessionStartInputHarness, SessionStartInputTrigger, TextBodyPart, TextPart, TextParts, +}; +use agentkit_core::{FinishReason, Item, ItemKind, Part, ToolOutput, TurnCancellation}; +use agentkit_loop::LoopError; + +pub(crate) use model::{HookedAdapter, HookedSession}; +pub(crate) use prompt::{PromptHook, PromptRecord}; +pub(crate) use tool::hooked_tools; + +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, +} + +impl Hooks { + /// Loads `~/.agents/hooks.json`, then `/.agents/hooks.json`. Project + /// hooks replace global hooks with the same id. Invalid files are skipped. + pub(crate) fn load(root: &Path) -> Option> { + let mut merged: Option = None; + for directory in crate::runtime::agents_directories(root).iter().rev() { + let path = directory.join(CONFIG_FILE); + match read(&path) { + Ok(Some(registration)) => merge(&mut merged, registration), + Ok(None) => {} + Err(error) => eprintln!("kit: ignoring {}: {error}", path.display()), + } + } + let registration = merged?; + match options().and_then(|options| Self::new(registration, options)) { + Ok(hooks) => Some(Arc::new(hooks)), + Err(error) => { + eprintln!("kit: agent hooks disabled: {error}"); + None + } + } + } + + fn new(registration: Registration, options: HooksOptions) -> Result { + let sdk = agenthooksprotocol::Hooks::new(registration, options) + .map_err(|error| error.to_string())?; + Ok(Self { sdk }) + } + + pub(crate) async fn session_start(&self, session: &str, trigger: SessionStartInputTrigger) { + let input = SessionStartInput::new( + SessionStartInputHarness::new("kit", env!("CARGO_PKG_VERSION")), + Vec::new(), + "allow-all".into(), + Box::new(Session::new(session)), + trigger, + ); + settle(EventType::SessionStart, self.sdk.session_start(input).await); + } + + pub(crate) async fn session_end(&self, session: &str, outcome: SessionEndInputOutcome) { + 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); + } + + /// Runs `context.compact.before`; a hook may supply the summary or replace + /// the summarizer instructions. + pub(crate) async fn compact_before( + &self, + items: &[Item], + manual: bool, + instructions: &str, + cancellation: Option<&TurnCancellation>, + ) -> Option { + let trigger = if manual { + ContextCompactBeforeInputTrigger::Manual + } else { + ContextCompactBeforeInputTrigger::Auto + }; + let mut input = ContextCompactBeforeInput::new(Vec::new(), trigger); + input.instructions = Some(Box::new(text_parts("instructions", instructions))); + let boundary = self + .sdk + .context_compact_before(input.with_sources().with_items(messages(items))); + let result = match cancellation { + Some(cancellation) => boundary.cancel_when(cancellation.cancelled()).await, + None => boundary.await, + }; + let outcome = settle(EventType::ContextCompactBefore, result)?; + if let Some(SettledCandidate::Summary(summary)) = outcome.outcome.candidate() { + return Some(Compaction::Summary(parts_text(summary))); + } + match outcome.settled_event { + Event::ContextCompactBefore(event) => match &event.0.instructions { + Presence::Present(edited) if parts_text(edited) != instructions => { + Some(Compaction::Instructions(parts_text(edited))) + } + _ => None, + }, + _ => None, + } + } + + /// Runs `context.compact.after`; a hook may replace the summary before it + /// is installed. + pub(crate) async fn compact_after(&self, items: &[Item], summary: String) -> String { + let removed = items + .iter() + .enumerate() + .map(|(index, item)| ContextCompactAfterInputRemovedItem::new(item_id(item, index))) + .collect(); + let input = ContextCompactAfterInput::new( + Box::new(ExecutionEventExecution::Executed( + ExecutionEventExecutionExecuted::new(), + )), + removed, + Box::new(text_parts("summary", &summary)), + ); + match settle( + EventType::ContextCompactAfter, + self.sdk.context_compact_after(input).await, + ) + .map(|outcome| outcome.settled_event) + { + Some(Event::ContextCompactAfter(event)) => parts_text(&event.0.summary), + _ => summary, + } + } +} + +pub(crate) fn session_outcome(result: &Result) -> SessionEndInputOutcome { + match result { + Ok(_) => SessionEndInputOutcome::Completed, + Err(LoopError::Cancelled) => SessionEndInputOutcome::Cancelled, + Err(_) => SessionEndInputOutcome::Error, + } +} + +pub(crate) enum Compaction { + Summary(String), + Instructions(String), +} + +fn options() -> Result { + use ModifyOperation::{Merge, Replace}; + let declarations = (|| { + Ok::<_, capability::Error>([ + (EventType::SessionStart, capability::observe()), + (EventType::SessionEnd, capability::observe()), + ( + EventType::UserMessageInbound, + capability::intercept() + .deny() + .modify_prompt([Replace])? + .build()?, + ), + ( + EventType::ModelRequestBefore, + capability::intercept().deny().build()?, + ), + (EventType::ModelResponseAfter, capability::observe()), + (EventType::ModelError, capability::observe()), + ( + EventType::ToolBefore, + capability::intercept() + .allow() + .deny() + .r#return() + .modify_input([Replace, Merge])? + .build()?, + ), + (EventType::ToolAfter, capability::observe()), + ( + EventType::ContextCompactBefore, + capability::intercept() + .r#return() + .modify_instructions([Replace])? + .build()?, + ), + ( + EventType::ContextCompactAfter, + capability::intercept().modify_summary([Replace])?.build()?, + ), + ]) + })() + .map_err(|error| error.to_string())?; + HooksOptions::from_declarations(SOURCE, declarations).map_err(|error| error.to_string()) +} + +fn read(path: &Path) -> Result, String> { + match crate::config_files::read(path) { + Ok(bytes) => serde_json::from_slice(&bytes) + .map(Some) + .map_err(|error| error.to_string()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(error) => Err(error.to_string()), + } +} + +fn merge(merged: &mut Option, next: Registration) { + let Some(registration) = merged else { + *merged = Some(next); + return; + }; + for backend in next.hooks { + match registration + .hooks + .iter_mut() + .find(|existing| existing.id == backend.id) + { + Some(existing) => *existing = backend, + None => registration.hooks.push(backend), + } + } + registration.protocol_version = next.protocol_version; +} + +fn report(event: EventType, diagnostics: &[DeliveryDiagnostic]) { + for diagnostic in diagnostics { + tracing::warn!( + event = event.as_str(), + ?diagnostic, + "agent hook delivery failed" + ); + } +} + +/// Reports delivery diagnostics; a dispatch error is reported and yields `None`. +fn settle( + event: EventType, + result: Result, HookError>, +) -> Option> { + match result { + Ok(outcome) => { + report(event, &outcome.diagnostics); + Some(outcome) + } + Err(error) => { + tracing::warn!(event = event.as_str(), %error, "agent hook dispatch failed"); + None + } + } +} + +fn denial(outcome: &ConsumerProtocolOutcome) -> String { + outcome + .responses + .iter() + .filter_map(|response| match &response.result.effects { + Presence::Present(effects) => Some(effects), + Presence::Missing => None, + }) + .flatten() + .find_map(|effect| match effect.as_ref() { + Effect::Deny(deny) => Some(deny.reason.clone()), + _ => None, + }) + .unwrap_or_else(|| "denied by agent hook".into()) +} + +fn messages(items: &[Item]) -> Vec> { + items + .iter() + .enumerate() + .filter_map(|(index, item)| message(item, index)) + .collect() +} + +fn item_id(item: &Item, index: usize) -> String { + item.id + .as_ref() + .map_or_else(|| format!("item-{index}"), ToString::to_string) +} + +fn message(item: &Item, index: usize) -> Option> { + let id = item_id(item, index); + let parts: Vec<_> = item + .parts + .iter() + .filter_map(part_text) + .enumerate() + .map(|(index, text)| text_part(&format!("{id}/{index}"), text)) + .collect(); + (!parts.is_empty()).then(|| MessageInput::new(id, role(item.kind), parts)) +} + +fn text_part(id: &str, text: impl Into) -> PartInput { + PartInput::text(TextBodyPart::new(id, text)) +} + +fn role(kind: ItemKind) -> CanonicalMessageRole { + match kind { + ItemKind::System => CanonicalMessageRole::System, + ItemKind::Developer | ItemKind::Context => CanonicalMessageRole::Developer, + ItemKind::User | ItemKind::Notification => CanonicalMessageRole::User, + ItemKind::Assistant => CanonicalMessageRole::Assistant, + ItemKind::Tool => CanonicalMessageRole::Tool, + } +} + +fn part_text(part: &Part) -> Option { + match part { + Part::Text(text) => Some(text.text.clone()), + Part::Structured(structured) => Some(structured.value.to_string()), + Part::ToolCall(call) => Some( + serde_json::Value::Object(serde_json::Map::from_iter([ + ("tool".to_owned(), call.name.clone().into()), + ("input".to_owned(), call.input.clone()), + ])) + .to_string(), + ), + Part::ToolResult(result) => output_text(&result.output), + Part::Media(_) | Part::File(_) | Part::Reasoning(_) | Part::Custom(_) => None, + } +} + +fn output_text(output: &ToolOutput) -> Option { + match output { + ToolOutput::Text(text) => Some(text.clone()), + ToolOutput::Structured(value) => Some(value.to_string()), + ToolOutput::Parts(parts) => { + let text: Vec<_> = parts.iter().filter_map(part_text).collect(); + (!text.is_empty()).then(|| text.join("\n")) + } + ToolOutput::Files(_) => None, + } +} + +fn finish_reason(reason: &FinishReason) -> String { + match reason { + FinishReason::Completed => "stop", + FinishReason::ToolCall => "tool_calls", + FinishReason::MaxTokens => "length", + FinishReason::Cancelled => "cancelled", + FinishReason::Blocked => "content_filter", + FinishReason::Error => "error", + FinishReason::Other(other) => other, + } + .to_owned() +} + +fn text(parts: &[Box]) -> String { + parts + .iter() + .filter_map(|part| match part.as_ref() { + ContentItem::TextBodyPart(part) => Some(part.text.as_str()), + _ => None, + }) + .collect() +} + +fn text_parts(id: &str, text: &str) -> TextParts { + TextParts::new(vec![Box::new(TextPart::from(TextBodyPart::new( + format!("{id}/0"), + text, + )))]) +} + +fn parts_text(parts: &TextParts) -> String { + parts + .0 + .iter() + .filter_map(|part| match part.as_ref() { + TextPart::TextBodyPart(part) => Some(part.text.as_str()), + _ => None, + }) + .collect() +} + +#[cfg(test)] +#[allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::disallowed_methods, + clippy::disallowed_macros +)] +pub(crate) mod test_support; +#[cfg(test)] +#[allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::disallowed_methods, + clippy::disallowed_macros +)] +mod tests; diff --git a/src/ahp/model.rs b/src/ahp/model.rs new file mode 100644 index 00000000..fc245fc3 --- /dev/null +++ b/src/ahp/model.rs @@ -0,0 +1,218 @@ +use std::sync::Arc; + +use agenthooksprotocol::inputs::{ + ModelErrorInput, ModelRequestBeforeInput, ModelResponseAfterInput, +}; +use agenthooksprotocol::{ + EventType, ExecutionEventAttempt, ExecutionEventError, ExecutionEventExecution, + ExecutionEventExecutionExecuted, ExecutionEventModel, Integer, ModelErrorInputExecution, + ModelErrorInputExecutionExecuted, ModelRequestBeforeInputParams, SessionStartInputTrigger, +}; +use agentkit_core::TurnCancellation; +use agentkit_loop::{ + LoopError, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, RetryObserver, SessionConfig, + TurnRequest, +}; +use async_trait::async_trait; + +use super::Hooks; + +/// Model adapter that reports session start and model calls to agent hooks. +#[derive(Clone)] +pub(crate) struct HookedAdapter { + inner: A, + hooks: Option>, + trigger: SessionStartInputTrigger, +} + +impl HookedAdapter { + pub(crate) fn new( + inner: A, + hooks: Option>, + trigger: SessionStartInputTrigger, + ) -> Self { + Self { + inner, + hooks, + trigger, + } + } +} + +#[async_trait] +impl ModelAdapter for HookedAdapter { + type Session = HookedSession; + + async fn start_session(&self, config: SessionConfig) -> Result { + let session_id = config.session_id.to_string(); + let inner = self.inner.start_session(config).await?; + if let Some(hooks) = &self.hooks { + hooks.session_start(&session_id, self.trigger.clone()).await; + } + Ok(HookedSession { + inner, + hooks: self.hooks.clone(), + attempts: 0, + }) + } + + fn provider_name(&self) -> Option<&str> { + self.inner.provider_name() + } +} + +pub(crate) struct HookedSession { + inner: S, + hooks: Option>, + attempts: u64, +} + +#[derive(Clone)] +struct Call { + hooks: Arc, + attempt: ExecutionEventAttempt, + model: ExecutionEventModel, +} + +impl Call { + async fn failed(&self, error: &LoopError) { + if matches!(error, LoopError::Cancelled) { + return; + } + let class = match error { + LoopError::ProviderFailure(_) | LoopError::Provider(_) => "provider", + _ => "adapter", + }; + let input = ModelErrorInput::new( + Box::new(self.attempt.clone()), + Box::new(ExecutionEventError::new(class, error.to_string())), + ModelErrorInputExecution::Executed(ModelErrorInputExecutionExecuted::new()), + Box::new(self.model.clone()), + ); + super::settle( + EventType::ModelError, + self.hooks.sdk.model_error(input).await, + ); + } +} + +#[async_trait] +impl ModelSession for HookedSession { + type Turn = HookedTurn; + + fn set_retry_observer(&mut self, observer: Option>) { + self.inner.set_retry_observer(observer); + } + + async fn begin_turn( + &mut self, + request: TurnRequest, + cancellation: Option, + ) -> Result { + let Some(hooks) = &self.hooks else { + let inner = self.inner.begin_turn(request, cancellation).await?; + return Ok(HookedTurn { inner, call: None }); + }; + self.attempts += 1; + let Some(number) = Integer::new(self.attempts.into()) else { + let inner = self.inner.begin_turn(request, cancellation).await?; + return Ok(HookedTurn { inner, call: None }); + }; + let call = Call { + hooks: Arc::clone(hooks), + attempt: ExecutionEventAttempt::new( + format!("{}/{}", request.turn_id, self.attempts), + number, + ), + model: ExecutionEventModel::new( + self.inner.model_name().unwrap_or("unknown"), + self.inner.provider_name().unwrap_or("unknown"), + ), + }; + let input = ModelRequestBeforeInput::new( + Box::new(call.attempt.clone()), + Vec::new(), + Box::new(call.model.clone()), + ModelRequestBeforeInputParams::new(), + ) + .with_sources() + .with_items(super::messages(&request.transcript)); + let boundary = hooks.sdk.model_request_before(input); + let result = match &cancellation { + Some(cancellation) => boundary.cancel_when(cancellation.cancelled()).await, + None => boundary.await, + }; + let reason = match super::settle(EventType::ModelRequestBefore, result) { + Some(outcome) if !outcome.outcome.is_denied() => None, + Some(outcome) => Some(super::denial(&outcome.outcome)), + None => Some("agent hook dispatch failed".into()), + }; + if let Some(reason) = reason { + return Err(LoopError::InvalidState(format!( + "model request denied by agent hook: {reason}" + ))); + } + match self.inner.begin_turn(request, cancellation).await { + Ok(inner) => Ok(HookedTurn { + inner, + call: Some(call), + }), + Err(error) => { + call.failed(&error).await; + Err(error) + } + } + } + + fn model_name(&self) -> Option<&str> { + self.inner.model_name() + } + + fn provider_name(&self) -> Option<&str> { + self.inner.provider_name() + } +} + +pub(crate) struct HookedTurn { + inner: T, + call: Option, +} + +#[async_trait] +impl ModelTurn for HookedTurn { + fn on_cancelled(&mut self) { + self.inner.on_cancelled(); + } + + async fn next_event( + &mut self, + cancellation: Option, + ) -> Result, LoopError> { + let event = self.inner.next_event(cancellation).await; + let Some(call) = &self.call else { + return event; + }; + match &event { + Ok(Some(ModelTurnEvent::Finished(result))) => { + let input = ModelResponseAfterInput::new( + Box::new(call.attempt.clone()), + Box::new(ExecutionEventExecution::Executed( + ExecutionEventExecutionExecuted::new(), + )), + super::finish_reason(&result.finish_reason), + Vec::new(), + Box::new(call.model.clone()), + ) + .with_sources() + .with_items(super::messages(&result.output_items)); + super::settle( + EventType::ModelResponseAfter, + call.hooks.sdk.model_response_after(input).await, + ); + } + Err(error) => call.failed(error).await, + _ => {} + } + event + } +} diff --git a/src/ahp/prompt.rs b/src/ahp/prompt.rs new file mode 100644 index 00000000..0fd7abb4 --- /dev/null +++ b/src/ahp/prompt.rs @@ -0,0 +1,145 @@ +use std::sync::Arc; + +use agenthooksprotocol::generated::Event; +use agenthooksprotocol::inputs::{MessageInput, UserMessageInboundInput}; +use agenthooksprotocol::{ + CanonicalMessageRole, EventType, Session, UserMessageInboundInputMessage, +}; +use agentkit_core::{ItemKind, Part}; +use agentkit_loop::{AgentEvent, LoopCtx, LoopError, LoopMutator, MutationPoint, TranscriptCursor}; +use async_trait::async_trait; + +use super::Hooks; +use crate::session::{SessionObserver, TransientTranscript}; + +const ACP_RESOURCE: &str = "acp.resource.uri"; + +/// Runs `user.message.inbound` over the prompt a new turn just admitted. +/// +/// Removing a denied prompt leaves no pending input, so the loop finishes the +/// turn without calling the model. +pub(crate) struct PromptHook { + hooks: Arc, + channel: &'static str, + record: PromptRecord, +} + +/// Where the transcript a prompt edit or denial rewrites is recorded. +pub(crate) enum PromptRecord { + Session(SessionObserver), + Transient(TransientTranscript), +} + +impl PromptRecord { + fn replace(&self, transcript: &[agentkit_core::Item]) -> Result<(), String> { + match self { + Self::Session(observer) => observer.replace(transcript), + Self::Transient(recorded) => { + recorded.replace(transcript); + Ok(()) + } + } + } +} + +impl PromptHook { + pub(crate) fn new(hooks: Arc, channel: &'static str, record: PromptRecord) -> Self { + Self { + hooks, + channel, + record, + } + } +} + +#[async_trait] +impl LoopMutator for PromptHook { + async fn mutate( + &self, + cursor: &mut TranscriptCursor<'_>, + ctx: LoopCtx<'_>, + ) -> Result<(), LoopError> { + if ctx.point != MutationPoint::AfterTurnEnded { + return Ok(()); + } + let Some(prompt) = cursor.last().filter(|item| item.kind == ItemKind::User) else { + return Ok(()); + }; + let index = cursor.len().saturating_sub(1); + let message = super::message(prompt, index).unwrap_or_else(|| { + MessageInput::new( + super::item_id(prompt, index), + CanonicalMessageRole::User, + Vec::new(), + ) + }); + let mut input = UserMessageInboundInput::new(UserMessageInboundInputMessage::new( + self.channel, + Vec::new(), + "user", + )); + input.session = Some(Box::new(Session::new(ctx.session_id.to_string()))); + let boundary = self + .hooks + .sdk + .user_message_inbound(input.with_sources().with_message_messages(vec![message])); + let result = match &ctx.cancellation { + Some(cancellation) => boundary.cancel_when(cancellation.cancelled()).await, + None => boundary.await, + }; + let edited = match super::settle(EventType::UserMessageInbound, result) { + Some(outcome) if outcome.outcome.is_denied() => Err(super::denial(&outcome.outcome)), + Some(outcome) => match outcome.settled_event { + Event::UserMessageInbound(event) => Ok(event + .message + .messages + .iter() + .map(|message| super::text(&message.parts)) + .collect::>() + .join("\n")), + _ => return Ok(()), + }, + None => Err("agent hook dispatch failed".into()), + }; + if ctx + .cancellation + .as_ref() + .is_some_and(|cancellation| cancellation.is_cancelled()) + { + return Err(LoopError::Cancelled); + } + match edited { + Ok(text) => { + let original = prompt + .parts + .iter() + .filter_map(super::part_text) + .collect::(); + if text == original { + return Ok(()); + } + if let Some(prompt) = cursor.last_mut() { + prompt.parts.retain(|part| super::part_text(part).is_none()); + prompt.parts.insert(0, Part::text(text)); + } + } + Err(reason) => { + cursor.pop(); + while cursor.last().is_some_and(|item| { + item.kind == ItemKind::Context && item.metadata.contains_key(ACP_RESOURCE) + }) { + cursor.pop(); + } + ctx.emitter.emit(AgentEvent::Warning { + message: format!("prompt blocked by agent hook: {reason}"), + }); + } + } + if let Err(error) = self.record.replace(cursor) { + ctx.emitter.emit(AgentEvent::Warning { + message: format!("could not persist agent hook prompt edit: {error}"), + }); + } + Ok(()) + } +} diff --git a/src/ahp/test_support.rs b/src/ahp/test_support.rs new file mode 100644 index 00000000..36cb0474 --- /dev/null +++ b/src/ahp/test_support.rs @@ -0,0 +1,105 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use agenthooksprotocol::adapters::registered::ManagedBackend; +use agenthooksprotocol::client::{HookError, LocalFuture}; +use serde_json::{Value, json}; + +use super::{Hooks, options}; + +const BACKEND: &str = "dev.kit.test"; + +/// In-process hook backend answering each event type with fixed effects. +#[derive(Default)] +pub(crate) struct Backend { + effects: Mutex>, + events: Mutex>, +} + +impl Backend { + pub(crate) fn respond(&self, event: &str, effects: Value) { + self.effects.lock().unwrap().insert(event.into(), effects); + } + + pub(crate) fn events(&self, event: &str) -> Vec { + self.events + .lock() + .unwrap() + .iter() + .filter(|delivered| delivered["type"] == event) + .cloned() + .collect() + } +} + +impl ManagedBackend for Backend { + fn call(&self, request: Value, _: Duration) -> LocalFuture<'_, Result> { + Box::pin(async move { + let event = request["params"]["event"].clone(); + self.events.lock().unwrap().push(event.clone()); + if request["method"] == "hooks/observe" { + return Ok(Value::Null); + } + let effects = event["type"] + .as_str() + .and_then(|name| self.effects.lock().unwrap().get(name).cloned()) + .unwrap_or_else(|| json!([])); + Ok(json!({ + "jsonrpc": "2.0", + "id": request["id"], + "result": {"protocolVersion": "draft", "effects": effects}, + })) + }) + } + + fn shutdown(&self) -> LocalFuture<'_, Result<(), HookError>> { + Box::pin(async { Ok(()) }) + } +} + +pub(crate) fn intercept(events: &[&str]) -> Value { + json!({ + "events": events, + "mode": "intercept", + "timeoutMs": 1000, + "failurePolicy": "fail-closed", + "content": {"default": "body"}, + }) +} + +pub(crate) fn observe(events: &[&str]) -> Value { + json!({"events": events, "mode": "observe", "content": {"default": "body"}}) +} + +pub(crate) fn hooks(subscriptions: Vec) -> (Arc, Arc) { + let backend = Arc::new(Backend::default()); + let registration = serde_json::from_value(json!({ + "protocolVersion": "draft", + "hooks": [{ + "id": BACKEND, + "transport": {"type": "stdio", "command": "unused", "lifecycle": "persistent"}, + "subscriptions": subscriptions, + }], + })) + .unwrap(); + let options = options().unwrap().with_backend(BACKEND, backend.clone()); + ( + Arc::new(Hooks::new(registration, options).unwrap()), + backend, + ) +} + +pub(crate) fn text_message(text: &str) -> Value { + json!([{ + "id": "edited", + "role": "user", + "parts": [{ + "id": "edited/0", + "kind": "text", + "mediaType": "text/plain", + "selection": "body", + "text": text, + }], + }]) +} diff --git a/src/ahp/tests.rs b/src/ahp/tests.rs new file mode 100644 index 00000000..b0faea3c --- /dev/null +++ b/src/ahp/tests.rs @@ -0,0 +1,576 @@ +use std::sync::{Arc, Mutex}; + +use agenthooksprotocol::SessionStartInputTrigger; +use agentkit_core::{ + DataRef, FinishReason, Item, ItemKind, MetadataMap, Modality, Part, SessionId, ToolCallId, + ToolOutput, ToolResultPart, TurnCancellation, TurnId, +}; +use agentkit_loop::{ + Agent, LoopError, LoopStep, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, + ModelTurnResult, SessionConfig, TurnRequest, +}; +use agentkit_tools_core::{ + AllowAllPermissions, OwnedToolContext, Tool, ToolContext, ToolError, ToolExecutionOutcome, + ToolName, ToolRegistry, ToolRequest, ToolResult, ToolSpec, +}; +use async_trait::async_trait; +use serde_json::{Value, json}; + +use super::test_support::{hooks, intercept, observe, text_message}; +use super::*; + +#[test] +fn every_granted_event_accepts_wildcard_subscriptions() { + hooks(vec![intercept(&["*"]), observe(&["*"])]); +} + +#[test] +fn project_hooks_replace_global_hooks_with_the_same_id() { + let registration = |hooks: Value| -> Registration { + serde_json::from_value(json!({"protocolVersion": "draft", "hooks": hooks})).unwrap() + }; + let backend = |id: &str, command: &str| { + json!({ + "id": id, + "transport": {"type": "stdio", "command": command, "lifecycle": "persistent"}, + "subscriptions": [observe(&["session.start"])], + }) + }; + let mut merged = None; + merge( + &mut merged, + registration(json!([ + backend("dev.kit.a", "global-a"), + backend("dev.kit.b", "global-b") + ])), + ); + merge( + &mut merged, + registration(json!([ + backend("dev.kit.b", "project-b"), + backend("dev.kit.c", "project-c") + ])), + ); + + let merged = serde_json::to_value(merged.unwrap()).unwrap(); + let commands: Vec<_> = merged["hooks"] + .as_array() + .unwrap() + .iter() + .map(|hook| hook["transport"]["command"].as_str().unwrap()) + .collect(); + assert_eq!(commands, ["global-a", "project-b", "project-c"]); +} + +#[test] +fn unreadable_registration_is_reported_and_missing_is_ignored() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join(CONFIG_FILE); + assert!(read(&path).unwrap().is_none()); + + crate::resilient_fs::write(&path, b"{\"hooks\": 1}").unwrap(); + assert!(read(&path).is_err()); +} + +struct Echo { + spec: ToolSpec, + current: Mutex>, + received: Mutex>, +} + +impl Echo { + fn new() -> Arc { + Arc::new(Self { + spec: ToolSpec::new( + ToolName::new("echo"), + "echo", + json!({ + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], + "additionalProperties": false, + }), + ), + current: Mutex::default(), + received: Mutex::default(), + }) + } +} + +#[async_trait] +impl Tool for Echo { + fn spec(&self) -> &ToolSpec { + &self.spec + } + + fn current_spec(&self) -> Option { + self.current.lock().unwrap().clone() + } + + async fn invoke_outcome( + &self, + request: ToolRequest, + context: &mut ToolContext<'_>, + ) -> ToolExecutionOutcome { + if request.input["text"] == "unavailable" { + return ToolExecutionOutcome::FailedBeforeInvocation(ToolError::Unavailable( + "offline".into(), + )); + } + match self.invoke(request, context).await { + Ok(result) => ToolExecutionOutcome::Completed(result), + Err(error) => ToolExecutionOutcome::Failed(error), + } + } + + async fn invoke( + &self, + request: ToolRequest, + _: &mut ToolContext<'_>, + ) -> Result { + self.received.lock().unwrap().push(request.input.clone()); + let text = request.input["text"] + .as_str() + .unwrap_or_default() + .to_owned(); + Ok(ToolResult::new(ToolResultPart::success( + request.call_id, + ToolOutput::text(text), + ))) + } +} + +async fn call(hooks: &Arc, echo: &Arc, input: Value) -> ToolExecutionOutcome { + let mut registry = ToolRegistry::new(); + registry.register_arc(echo.clone()); + let tool = hooked_tools(registry, hooks) + .get(&ToolName::new("echo")) + .unwrap(); + let context = OwnedToolContext { + session_id: SessionId::new("session"), + turn_id: TurnId::new("turn"), + metadata: MetadataMap::new(), + permissions: Arc::new(AllowAllPermissions), + resources: Arc::new(()), + cancellation: None, + execution_scope: None, + approved_request: None, + }; + let request = ToolRequest::new( + ToolCallId::new("call"), + ToolName::new("echo"), + input, + context.session_id.clone(), + context.turn_id.clone(), + ); + tool.invoke_outcome(request, &mut context.borrowed()).await +} + +fn output(outcome: ToolExecutionOutcome) -> ToolOutput { + match outcome { + ToolExecutionOutcome::Completed(result) => result.result.output, + other => panic!("expected a completed call, got {other:?}"), + } +} + +fn denial_message(outcome: ToolExecutionOutcome) -> String { + match outcome { + ToolExecutionOutcome::FailedBeforeInvocation(ToolError::PermissionDenied(denial)) => { + denial.message + } + other => panic!("expected a denied call, got {other:?}"), + } +} + +#[tokio::test] +async fn tool_before_denial_skips_the_tool() { + let (hooks, backend) = hooks(vec![intercept(&["tool.before"])]); + backend.respond( + "tool.before", + json!([{"type": "deny", "reason": "not here"}]), + ); + let echo = Echo::new(); + + let outcome = call(&hooks, &echo, json!({"text": "hello"})).await; + + assert_eq!(denial_message(outcome), "not here"); + assert!(echo.received.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn tool_before_input_edits_reach_the_tool() { + let (hooks, backend) = hooks(vec![intercept(&["tool.before"])]); + backend.respond( + "tool.before", + json!([{"type": "modify", "target": "input", "operation": "merge", "value": {"text": "edited"}}]), + ); + let echo = Echo::new(); + + let outcome = call(&hooks, &echo, json!({"text": "hello"})).await; + + assert_eq!(output(outcome), ToolOutput::text("edited")); + assert_eq!(*echo.received.lock().unwrap(), [json!({"text": "edited"})]); +} + +#[tokio::test] +async fn tool_before_edits_must_satisfy_the_tool_schema() { + let (hooks, backend) = hooks(vec![intercept(&["tool.before"])]); + backend.respond( + "tool.before", + json!([{"type": "modify", "target": "input", "operation": "merge", "value": {"extra": true}}]), + ); + let echo = Echo::new(); + + let outcome = call(&hooks, &echo, json!({"text": "hello"})).await; + + denial_message(outcome); + assert!(echo.received.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn tool_before_validates_against_the_current_schema() { + let (hooks, backend) = hooks(vec![intercept(&["tool.before"])]); + backend.respond( + "tool.before", + json!([{"type": "modify", "target": "input", "operation": "merge", "value": {"extra": true}}]), + ); + let echo = Echo::new(); + denial_message(call(&hooks, &echo, json!({"text": "hello"})).await); + + *echo.current.lock().unwrap() = Some(ToolSpec::new( + ToolName::new("echo"), + "echo", + json!({"type": "object"}), + )); + call(&hooks, &echo, json!({"text": "hello"})).await; + + assert_eq!( + *echo.received.lock().unwrap(), + [json!({"text": "hello", "extra": true})] + ); +} + +#[tokio::test] +async fn unsubscribed_events_leave_tools_and_prompts_unchanged() { + let (hooks, _) = hooks(vec![observe(&["session.start"])]); + let echo = Echo::new(); + let model = Model::default(); + + let outcome = call(&hooks, &echo, json!({"text": "hello"})).await; + let (_, persisted) = prompt(&hooks, &model, "original").await; + + assert_eq!(output(outcome), ToolOutput::text("hello")); + let requests = model.requests.lock().unwrap(); + assert_eq!(last_user_text(&requests[0]).as_deref(), Some("original")); + assert_eq!(last_user_text(&persisted).as_deref(), Some("original")); +} + +#[tokio::test] +async fn tool_before_return_supplies_the_result() { + let (hooks, backend) = hooks(vec![intercept(&["tool.before"])]); + backend.respond( + "tool.before", + json!([{"type": "return", "value": "cached"}]), + ); + let echo = Echo::new(); + + let outcome = call(&hooks, &echo, json!({"text": "hello"})).await; + + assert_eq!(output(outcome), ToolOutput::text("cached")); + assert!(echo.received.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn tool_after_observes_the_output() { + let (hooks, backend) = hooks(vec![observe(&["tool.after"])]); + let echo = Echo::new(); + + call(&hooks, &echo, json!({"text": "hello"})).await; + + let events = backend.events("tool.after"); + assert_eq!(events.len(), 1); + assert_eq!(events[0]["outcome"], "ok"); + assert_eq!(events[0]["items"][0]["parts"][0]["text"], "hello"); +} + +#[tokio::test] +async fn tool_after_reports_hook_supplied_results_as_skipped() { + let (hooks, backend) = hooks(vec![intercept(&["tool.before"]), observe(&["tool.after"])]); + backend.respond( + "tool.before", + json!([{"type": "return", "value": "cached"}]), + ); + let echo = Echo::new(); + + call(&hooks, &echo, json!({"text": "hello"})).await; + + let events = backend.events("tool.after"); + assert_eq!(events.len(), 1); + assert_eq!(events[0]["execution"]["status"], "skipped"); + assert_eq!(events[0]["execution"]["reason"], "supplied_result"); + assert_eq!(events[0]["items"][0]["parts"][0]["text"], "cached"); +} + +#[tokio::test] +async fn tool_after_reports_pre_invocation_failures_as_skipped() { + let (hooks, backend) = hooks(vec![observe(&["tool.after"])]); + let echo = Echo::new(); + + call(&hooks, &echo, json!({"text": "unavailable"})).await; + + let events = backend.events("tool.after"); + assert_eq!(events[0]["outcome"], "error"); + assert_eq!(events[0]["execution"]["status"], "skipped"); + assert_eq!(events[0]["execution"]["reason"], "other"); +} + +/// Answers every turn with "done" and records the transcripts it was sent. +#[derive(Clone, Default)] +struct Model { + requests: Arc>>>, +} + +#[async_trait] +impl ModelAdapter for Model { + type Session = Self; + + async fn start_session(&self, _: SessionConfig) -> Result { + Ok(self.clone()) + } +} + +#[async_trait] +impl ModelSession for Model { + type Turn = Reply; + + async fn begin_turn( + &mut self, + request: TurnRequest, + _: Option, + ) -> Result { + self.requests.lock().unwrap().push(request.transcript); + Ok(Reply(Some(ModelTurnResult { + finish_reason: FinishReason::Completed, + output_items: vec![Item::text(ItemKind::Assistant, "done")], + usage: None, + metadata: MetadataMap::new(), + model: None, + response_id: None, + }))) + } + + fn model_name(&self) -> Option<&str> { + Some("model") + } + + fn provider_name(&self) -> Option<&str> { + Some("provider") + } +} + +struct Reply(Option); + +#[async_trait] +impl ModelTurn for Reply { + async fn next_event( + &mut self, + _: Option, + ) -> Result, LoopError> { + Ok(self.0.take().map(ModelTurnEvent::Finished)) + } +} + +fn last_user_text(transcript: &[Item]) -> Option { + transcript + .iter() + .rev() + .find(|item| item.kind == ItemKind::User) + .map(|item| item.parts.iter().filter_map(part_text).collect()) +} + +async fn prompt(hooks: &Arc, model: &Model, text: &str) -> (LoopStep, Vec) { + admit(hooks, model, vec![Item::text(ItemKind::User, text)]).await +} + +async fn admit(hooks: &Arc, model: &Model, input: Vec) -> (LoopStep, Vec) { + let root = tempfile::tempdir().unwrap(); + let directory = root.path().join("sessions"); + let initial = vec![Item::text(ItemKind::System, "system")]; + let opened = + crate::session::open_in(root.path(), &directory, "abc", false, false, initial).unwrap(); + let mut driver = Agent::builder() + .model(HookedAdapter::new( + model.clone(), + Some(hooks.clone()), + SessionStartInputTrigger::Startup, + )) + .mutator(PromptHook::new( + hooks.clone(), + "cli", + PromptRecord::Session(opened.observer.clone()), + )) + .transcript_observer(opened.observer) + .transcript(opened.transcript) + .input(input) + .build() + .unwrap() + .start(SessionConfig::new("abc")) + .await + .unwrap(); + let step = driver.next().await.unwrap(); + let persisted = crate::session::load_in(root.path(), &directory, "abc").unwrap(); + (step, persisted) +} + +#[tokio::test] +async fn prompt_edits_replace_the_submitted_text() { + let (hooks, backend) = hooks(vec![intercept(&["user.message.inbound"])]); + backend.respond( + "user.message.inbound", + json!([{"type": "modify", "target": "prompt", "operation": "replace", "value": text_message("edited")}]), + ); + let model = Model::default(); + + let (_, persisted) = prompt(&hooks, &model, "original").await; + + let requests = model.requests.lock().unwrap(); + assert_eq!(last_user_text(&requests[0]).as_deref(), Some("edited")); + assert_eq!(last_user_text(&persisted).as_deref(), Some("edited")); +} + +#[tokio::test] +async fn denied_prompts_finish_without_a_model_call() { + let (hooks, backend) = hooks(vec![intercept(&["user.message.inbound"])]); + backend.respond( + "user.message.inbound", + json!([{"type": "deny", "reason": "blocked"}]), + ); + let model = Model::default(); + + let (step, persisted) = prompt(&hooks, &model, "original").await; + + assert!(matches!(step, LoopStep::Finished(result) if result.items.is_empty())); + assert!(model.requests.lock().unwrap().is_empty()); + assert_eq!(last_user_text(&persisted), None); +} + +#[tokio::test] +async fn media_only_prompts_reach_the_inbound_hook() { + let (hooks, backend) = hooks(vec![intercept(&["user.message.inbound"])]); + backend.respond( + "user.message.inbound", + json!([{"type": "deny", "reason": "blocked"}]), + ); + let model = Model::default(); + let image = Item::new( + ItemKind::User, + vec![Part::media( + Modality::Image, + "image/png", + DataRef::uri("file:///image.png"), + )], + ); + + let (step, _) = admit(&hooks, &model, vec![image]).await; + + assert!(matches!(step, LoopStep::Finished(result) if result.items.is_empty())); + assert!(model.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn denied_prompts_drop_their_attached_resources() { + let (hooks, backend) = hooks(vec![intercept(&["user.message.inbound"])]); + backend.respond( + "user.message.inbound", + json!([{"type": "deny", "reason": "blocked"}]), + ); + let model = Model::default(); + let mut resource = Item::text(ItemKind::Context, "attached secret"); + resource + .metadata + .insert("acp.resource.uri".into(), json!("file:///secret.txt")); + + let (_, persisted) = admit( + &hooks, + &model, + vec![resource, Item::text(ItemKind::User, "original")], + ) + .await; + + assert!(persisted.iter().all(|item| item.kind == ItemKind::System)); +} + +#[tokio::test] +async fn model_request_denial_fails_the_turn() { + let (hooks, backend) = hooks(vec![intercept(&["model.request.before"])]); + backend.respond( + "model.request.before", + json!([{"type": "deny", "reason": "budget"}]), + ); + let model = Model::default(); + let mut session = HookedAdapter::new( + model.clone(), + Some(hooks), + SessionStartInputTrigger::Startup, + ) + .start_session(SessionConfig::new("abc")) + .await + .unwrap(); + + let request = TurnRequest { + session_id: SessionId::new("abc"), + turn_id: TurnId::new("turn"), + transcript: vec![Item::text(ItemKind::User, "hello")], + available_tools: Vec::new(), + cache: None, + metadata: MetadataMap::new(), + }; + let error = session.begin_turn(request, None).await.err().unwrap(); + + assert!(error.to_string().contains("budget"), "{error}"); + assert!(model.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn session_start_and_model_response_are_observed() { + let (hooks, backend) = hooks(vec![observe(&["session.start", "model.response.after"])]); + let model = Model::default(); + + prompt(&hooks, &model, "hello").await; + + let started = backend.events("session.start"); + assert_eq!(started.len(), 1); + assert_eq!(started[0]["session"]["id"], "abc"); + assert_eq!(started[0]["trigger"], "startup"); + let responses = backend.events("model.response.after"); + assert_eq!(responses.len(), 1); + assert_eq!(responses[0]["finishReason"], "stop"); + assert_eq!(responses[0]["items"][0]["parts"][0]["text"], "done"); +} + +#[tokio::test] +async fn compaction_hooks_supply_and_rewrite_the_summary() { + let (hooks, backend) = hooks(vec![intercept(&[ + "context.compact.before", + "context.compact.after", + ])]); + backend.respond( + "context.compact.before", + json!([{"type": "return", "value": [{"id": "s/0", "kind": "text", "mediaType": "text/plain", "selection": "body", "text": "supplied"}]}]), + ); + let items = [Item::text(ItemKind::User, "old")]; + + let before = hooks + .compact_before(&items, true, "instructions", None) + .await; + assert!(matches!(before, Some(Compaction::Summary(summary)) if summary == "supplied")); + + backend.respond( + "context.compact.after", + json!([{"type": "modify", "target": "summary", "operation": "replace", "value": [{"id": "s/0", "kind": "text", "mediaType": "text/plain", "selection": "body", "text": "rewritten"}]}]), + ); + assert_eq!( + hooks.compact_after(&items, "draft".into()).await, + "rewritten" + ); +} diff --git a/src/ahp/tool.rs b/src/ahp/tool.rs new file mode 100644 index 00000000..d66cd6ec --- /dev/null +++ b/src/ahp/tool.rs @@ -0,0 +1,310 @@ +use std::sync::{Arc, Mutex}; + +use agenthooksprotocol::client::{Decision, HookError}; +use agenthooksprotocol::contracts::Contract; +use agenthooksprotocol::generated::Presence; +use agenthooksprotocol::inputs::{MessageInput, ToolAfterInput, ToolBeforeInput}; +use agenthooksprotocol::{ + Attachment, CanonicalMessageRole, EventType, ExecutionEventError, ExecutionEventExecution, + ExecutionEventExecutionExecuted, ExecutionEventExecutionSkippedCancelled, + ExecutionEventExecutionSkippedOther, ExecutionEventExecutionSkippedPolicy, + ExecutionEventExecutionSkippedSuppliedResult, ToolAfterInputOrigin, ToolAfterInputOutcome, + ToolBeforeInputOrigin, state, +}; +use agentkit_core::{ToolOutput, ToolResultPart}; +use agentkit_tools_core::{ + PermissionCode, PermissionDecision, PermissionDenial, PermissionRequest, Tool, ToolContext, + ToolError, ToolExecutionOutcome, ToolInterruption, ToolRegistry, ToolRequest, ToolResult, + ToolSpec, +}; +use async_trait::async_trait; +use serde_json::Value; + +use super::Hooks; + +/// Wraps every tool so `tool.before` and `tool.after` surround its invocation. +pub(crate) fn hooked_tools(registry: ToolRegistry, hooks: &Arc) -> ToolRegistry { + registry + .tools() + .into_iter() + .fold(ToolRegistry::new(), |mut hooked, tool| { + hooked.register(HookedTool { + tool, + hooks: Arc::clone(hooks), + validator: Mutex::default(), + }); + hooked + }) +} + +struct HookedTool { + tool: Arc, + hooks: Arc, + validator: Mutex)>>, +} + +enum Before { + Run(ToolRequest), + Supplied(ToolRequest, ToolExecutionOutcome), + Return(ToolExecutionOutcome), +} + +impl HookedTool { + async fn before(&self, mut request: ToolRequest, context: &ToolContext<'_>) -> Before { + let validator = match self.validator() { + Ok(validator) => validator, + Err(error) => return Before::Return(denied(error)), + }; + let input = ToolBeforeInput::new( + request.call_id.to_string(), + request.tool_name.to_string(), + request.input.clone(), + request.tool_name.to_string(), + ToolBeforeInputOrigin::Native, + ); + let boundary = match self + .hooks + .sdk + .tool_before(input) + .arguments_contract(Contract::::validated(move |input| { + validator.validate(input).map_err(|error| error.to_string()) + })) + .initial_snapshot(state::initial(Decision::Allow)) + { + Ok(boundary) => boundary, + Err(error) => return Before::Return(denied(error.to_string())), + }; + let result = match &context.cancellation { + Some(cancellation) => boundary.cancel_when(cancellation.cancelled()).await, + None => boundary.await, + }; + let outcome = match result { + Ok(outcome) => outcome, + Err(error) => { + tracing::warn!(event = "tool.before", %error, "agent hook dispatch failed"); + return Before::Return(denied(error.to_string())); + } + }; + super::report(EventType::ToolBefore, &outcome.diagnostics); + if outcome.outcome.is_denied() { + return Before::Return(denied(super::denial(&outcome.outcome))); + } + if let Some(value) = outcome.outcome.raw().supplied_result() { + let output = match value { + Value::String(text) => ToolOutput::text(text.clone()), + value => ToolOutput::structured(value.clone()), + }; + let result = ToolResultPart::success(request.call_id.clone(), output); + return Before::Supplied( + request, + ToolExecutionOutcome::Completed(ToolResult::new(result)), + ); + } + let reauthorize = !outcome.outcome.can_execute(); + if reauthorize + && !(outcome.outcome.raw().approval_invalidated && outcome.outcome.requires_approval()) + { + return Before::Return(denied("agent hook did not authorize the call".into())); + } + request.input = match outcome.input { + Ok(input) => input, + Err(error) => return Before::Return(denied(error.to_string())), + }; + if reauthorize && let Some(outcome) = self.authorize(&request, context) { + return Before::Return(outcome); + } + Before::Run(request) + } + + fn validator(&self) -> Result, String> { + let mut cached = self + .validator + .lock() + .map_err(|_| "tool schema cache is poisoned".to_owned())?; + let schema = match self.tool.current_spec() { + Some(spec) => spec.input_schema, + None => self.tool.spec().input_schema.clone(), + }; + if let Some((cached_schema, validator)) = cached.as_ref() + && *cached_schema == schema + { + return Ok(Arc::clone(validator)); + } + let validator = Arc::new( + jsonschema::validator_for(&schema) + .map_err(|error| format!("invalid tool schema: {error}"))?, + ); + *cached = Some((schema, Arc::clone(&validator))); + Ok(validator) + } + + /// Applies the host permission policy to input a hook edited; `None` + /// authorizes the call. + fn authorize( + &self, + request: &ToolRequest, + context: &ToolContext<'_>, + ) -> Option { + let requests = match self.tool.proposed_requests(request) { + Ok(requests) => requests, + Err(error) => return Some(ToolExecutionOutcome::FailedBeforeInvocation(error)), + }; + for permission in requests { + match context.permissions.evaluate(permission.as_ref()) { + PermissionDecision::Allow => {} + PermissionDecision::Deny(denial) => { + return Some(ToolExecutionOutcome::FailedBeforeInvocation( + ToolError::PermissionDenied(denial), + )); + } + PermissionDecision::RequireApproval(mut approval) => { + let approved = context + .approved_request + .as_ref() + .is_some_and(|approved| approved.id == approval.id); + if !approved { + approval.call_id = Some(request.call_id.clone()); + return Some(ToolExecutionOutcome::Interrupted( + ToolInterruption::ApprovalRequired(approval), + )); + } + } + } + } + None + } + + async fn after(&self, request: &ToolRequest, outcome: &ToolExecutionOutcome, supplied: bool) { + let execution = match outcome { + _ if supplied => ExecutionEventExecution::SkippedSuppliedResult( + ExecutionEventExecutionSkippedSuppliedResult::new(), + ), + ToolExecutionOutcome::FailedBeforeInvocation(ToolError::Cancelled) => { + ExecutionEventExecution::SkippedCancelled( + ExecutionEventExecutionSkippedCancelled::new(), + ) + } + ToolExecutionOutcome::FailedBeforeInvocation(ToolError::PermissionDenied(_)) => { + ExecutionEventExecution::SkippedPolicy(ExecutionEventExecutionSkippedPolicy::new()) + } + ToolExecutionOutcome::FailedBeforeInvocation(error) => { + let mut skipped = ExecutionEventExecutionSkippedOther::new(); + skipped.detail = Presence::Present(error.to_string()); + ExecutionEventExecution::SkippedOther(skipped) + } + _ => ExecutionEventExecution::Executed(ExecutionEventExecutionExecuted::new()), + }; + let (status, output, error) = match outcome { + ToolExecutionOutcome::Completed(result) if result.result.is_error => ( + ToolAfterInputOutcome::Error, + super::output_text(&result.result.output), + None, + ), + ToolExecutionOutcome::Completed(result) => ( + ToolAfterInputOutcome::Ok, + super::output_text(&result.result.output), + None, + ), + ToolExecutionOutcome::Failed(ToolError::Cancelled) + | ToolExecutionOutcome::FailedBeforeInvocation(ToolError::Cancelled) => { + (ToolAfterInputOutcome::Cancelled, None, None) + } + ToolExecutionOutcome::Failed(error) + | ToolExecutionOutcome::FailedBeforeInvocation(error) => { + (ToolAfterInputOutcome::Error, None, Some(error.to_string())) + } + ToolExecutionOutcome::Interrupted(_) => return, + }; + let mut input = ToolAfterInput::new( + request.call_id.to_string(), + Box::new(execution), + Vec::new(), + status, + request.tool_name.to_string(), + request.input.clone(), + request.tool_name.to_string(), + ToolAfterInputOrigin::Native, + ); + input.error = error.map(|message| Box::new(ExecutionEventError::new("tool", message))); + let items = output + .map(|text| { + vec![MessageInput::new( + format!("{}/output", request.call_id), + CanonicalMessageRole::Tool, + vec![super::text_part( + &format!("{}/output/0", request.call_id), + text, + )], + )] + }) + .unwrap_or_default(); + let result = match input + .with_sources::() + .with_items(items) + .into_host_event() + { + Ok((event, _)) => self.hooks.sdk.event_for("tool.after", event).await, + Err(error) => Err(HookError(error.to_string())), + }; + super::settle(EventType::ToolAfter, result); + } +} + +fn denied(message: String) -> ToolExecutionOutcome { + ToolExecutionOutcome::FailedBeforeInvocation(ToolError::PermissionDenied(PermissionDenial { + code: PermissionCode::CustomPolicyDenied, + message, + metadata: Default::default(), + })) +} + +#[async_trait] +impl Tool for HookedTool { + fn spec(&self) -> &ToolSpec { + self.tool.spec() + } + + fn current_spec(&self) -> Option { + self.tool.current_spec() + } + + fn proposed_requests( + &self, + request: &ToolRequest, + ) -> Result>, ToolError> { + self.tool.proposed_requests(request) + } + + async fn invoke( + &self, + request: ToolRequest, + context: &mut ToolContext<'_>, + ) -> Result { + match self.invoke_outcome(request, context).await { + ToolExecutionOutcome::Completed(result) => Ok(result), + ToolExecutionOutcome::Failed(error) + | ToolExecutionOutcome::FailedBeforeInvocation(error) => Err(error), + ToolExecutionOutcome::Interrupted(_) => Err(ToolError::Internal( + "tool interruption requires invoke_outcome".into(), + )), + } + } + + async fn invoke_outcome( + &self, + request: ToolRequest, + context: &mut ToolContext<'_>, + ) -> ToolExecutionOutcome { + let request = match self.before(request, context).await { + Before::Run(request) => request, + Before::Supplied(request, outcome) => { + self.after(&request, &outcome, true).await; + return outcome; + } + Before::Return(outcome) => return outcome, + }; + let outcome = self.tool.invoke_outcome(request.clone(), context).await; + self.after(&request, &outcome, false).await; + outcome + } +} diff --git a/src/compaction.rs b/src/compaction.rs index e868f10f..8f34ac40 100644 --- a/src/compaction.rs +++ b/src/compaction.rs @@ -13,7 +13,7 @@ use agentkit_loop::{ MutationPoint, SessionConfig, TelemetryConfig, TranscriptCursor, }; use async_trait::async_trait; -use std::{collections::HashMap, time::Instant}; +use std::{collections::HashMap, sync::Arc, time::Instant}; use crate::{ events::{self, RuntimeEvent}, @@ -373,6 +373,7 @@ pub fn automatic( telemetry: TelemetryConfig, persistence: Option, session_id: impl Into, + hooks: Option>, ) -> Result where M: ModelAdapter + Clone + 'static, @@ -381,6 +382,7 @@ where adapter, telemetry, session_id: session_id.into(), + hooks, }; let inner = StrategyCompactor::new( |transcript: &[Item], _point: MutationPoint| { @@ -406,6 +408,7 @@ struct KitCompactionBackend { adapter: M, telemetry: TelemetryConfig, session_id: SessionId, + hooks: Option>, } fn build_summary_prompt(request: &SummaryRequest) -> Result { @@ -476,11 +479,31 @@ where { return Err(CompactionError::Cancelled); } + let mut instructions = SUMMARY_PROMPT.to_owned(); + if let Some(hooks) = &self.hooks { + match hooks + .compact_before( + &request.items, + request.reason == CompactionReason::Manual, + SUMMARY_PROMPT, + cancellation.as_ref(), + ) + .await + { + Some(crate::ahp::Compaction::Summary(summary)) => { + return self.installed(&request, summary).await; + } + Some(crate::ahp::Compaction::Instructions(replacement)) => { + instructions = replacement; + } + None => {} + } + } let prompt = build_summary_prompt(&request)?; let mut builder = Agent::builder() .model(self.adapter.clone()) .input(vec![ - Item::text(ItemKind::System, SUMMARY_PROMPT), + Item::text(ItemKind::System, instructions), Item::text(ItemKind::User, prompt), ]) .telemetry(self.telemetry); @@ -520,9 +543,23 @@ where Err(error) => return Err(CompactionError::Failed(error.to_string())), } }; + self.installed(&request, summary).await + } +} + +impl KitCompactionBackend { + async fn installed( + &self, + request: &SummaryRequest, + summary: String, + ) -> Result { + let summary = match &self.hooks { + Some(hooks) => hooks.compact_after(&request.items, summary).await, + None => summary, + }; if summary.trim().is_empty() { return Err(CompactionError::Failed( - "compaction agent returned an empty summary".into(), + "compaction summary is empty".into(), )); } let mut metadata = MetadataMap::new(); @@ -844,6 +881,7 @@ mod tests { use super::*; + #[derive(Clone)] struct NoModelTurn; #[async_trait] @@ -878,6 +916,65 @@ mod tests { } } + #[tokio::test] + async fn hook_supplied_summary_skips_the_summarizer_model() { + use crate::ahp::test_support::{hooks, intercept}; + let (hooks, backend) = hooks(vec![intercept(&["context.compact.before"])]); + backend.respond( + "context.compact.before", + json!([{"type": "return", "value": [{"id": "s/0", "kind": "text", "mediaType": "text/plain", "selection": "body", "text": "from hook"}]}]), + ); + let compactor = KitCompactionBackend { + adapter: NoModelTurn, + telemetry: TelemetryConfig::default(), + session_id: SessionId::new("compaction"), + hooks: Some(hooks), + }; + + let result = compactor + .summarize( + SummaryRequest::new( + vec![Item::text(ItemKind::User, "old request")], + CompactionReason::Manual, + ), + None, + ) + .await + .unwrap(); + + assert_eq!(result.items.len(), 1); + assert!(is_compaction_summary(&result.items[0])); + assert_eq!(item_text(&result.items[0]), Some("from hook")); + } + + #[tokio::test] + async fn blank_hook_summary_fails_compaction() { + use crate::ahp::test_support::{hooks, intercept}; + let (hooks, backend) = hooks(vec![intercept(&["context.compact.before"])]); + backend.respond( + "context.compact.before", + json!([{"type": "return", "value": [{"id": "s/0", "kind": "text", "mediaType": "text/plain", "selection": "body", "text": " "}]}]), + ); + let compactor = KitCompactionBackend { + adapter: NoModelTurn, + telemetry: TelemetryConfig::default(), + session_id: SessionId::new("compaction"), + hooks: Some(hooks), + }; + + let result = compactor + .summarize( + SummaryRequest::new( + vec![Item::text(ItemKind::User, "old request")], + CompactionReason::Manual, + ), + None, + ) + .await; + + assert!(matches!(result, Err(CompactionError::Failed(_)))); + } + #[tokio::test] async fn manual_compaction_continues_after_persistence_loses_ownership() { let root = tempfile::tempdir().unwrap(); diff --git a/src/docs.rs b/src/docs.rs index 321915f0..52f8b8f7 100644 --- a/src/docs.rs +++ b/src/docs.rs @@ -414,6 +414,7 @@ mod tests { assert_eq!( paths, vec![ + "docs/user/agent-hooks-protocol.md", "docs/user/agent-plugins.md", "docs/user/compose-and-local-tools.md", "docs/user/evaluations.md", diff --git a/src/lib.rs b/src/lib.rs index 1c71d1e5..2fcf2437 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,6 +15,7 @@ )] mod acp_child; +mod ahp; mod artifacts; pub mod compaction; mod compose_output; diff --git a/src/main.rs b/src/main.rs index 652a9c15..731e6628 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2271,6 +2271,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { )?, }; let runtime = kit::Runtime::with_eval(runtime, config.experimental.eval)?; + let runtime = kit::Runtime::with_hooks(runtime)?; let runtime = kit::Runtime::with_plugin_runtime(runtime, plugins)?; let runtime = kit::Runtime::with_telemetry(runtime, telemetry_settings.clone())?; let (harnesses, default_harness) = config.harnesses()?; @@ -2369,6 +2370,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { )?, }; let runtime = kit::Runtime::with_eval(runtime, config.experimental.eval)?; + let runtime = kit::Runtime::with_hooks(runtime)?; let runtime = kit::Runtime::with_plugin_runtime(runtime, plugins)?; let runtime = kit::Runtime::with_telemetry(runtime, telemetry_settings.clone())?; let runtime = kit::Runtime::with_depth(runtime, subagent_depth)?; @@ -2437,6 +2439,7 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { openrouter_api_key.as_ref().map(|(key, _)| key.clone()), )?; let runtime = kit::Runtime::with_eval(runtime, config.experimental.eval)?; + let runtime = kit::Runtime::with_hooks(runtime)?; let runtime = kit::Runtime::with_plugin_runtime(runtime, plugins)?; let runtime = kit::Runtime::with_telemetry(runtime, telemetry_settings.clone())?; let (harnesses, default_harness) = config.harnesses()?; diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index 4ba0e85f..f1465e73 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -11,6 +11,7 @@ use std::{ }; use agent_client_protocol::{Client, ConnectTo, ConnectionTo, Handled}; +use agenthooksprotocol::SessionEndInputOutcome; use agentkit_acp::{ AcpClientHandle, AcpClientMessage, AcpIntegration, AcpRuntimeError, AcpSessionBinding, AgentAuthCapabilities, AudioContent, AuthMethod, AuthMethodTerminal, AuthenticateRequest, @@ -2096,7 +2097,8 @@ async fn session_actor(actor: SessionActor) { let _ = reply.send(result); } Some(Command::Close { reply }) => { - clean_up_session(&session_id, &mut driver, &tasks, &background_jobs).await; + clean_up_session(&runtime, &session_id, &mut driver, &tasks, &background_jobs) + .await; // A close acknowledgement means the actor-owned binding is // already gone, so callers can immediately reuse the id. drop(binding.take()); @@ -2104,7 +2106,8 @@ async fn session_actor(actor: SessionActor) { break; } None => { - clean_up_session(&session_id, &mut driver, &tasks, &background_jobs).await; + clean_up_session(&runtime, &session_id, &mut driver, &tasks, &background_jobs) + .await; break; } }, @@ -2327,6 +2330,7 @@ pub(super) async fn cancel_background_jobs( } pub(super) async fn clean_up_session( + runtime: &Runtime, session_id: &agentkit_acp::SessionId, driver: &mut LoopDriver, tasks: &TaskManagerHandle, @@ -2346,6 +2350,11 @@ pub(super) async fn clean_up_session( if let Err(error) = driver.cancel_pending_approvals().await { eprintln!("failed to cancel ACP approvals for {session_id}: {error}"); } + if let Some(hooks) = runtime.hooks() { + hooks + .session_end(&session_id.to_string(), SessionEndInputOutcome::Completed) + .await; + } } fn record_acp_runtime_failure( diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 658a1b66..24be0134 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -1472,14 +1472,28 @@ async fn session_actor(actor: SessionActor) } Some(Command::Close { reply }) => { let v1_id = agentkit_acp::SessionId::new(session_id.to_string()); - super::clean_up_session(&v1_id, &mut driver, &tasks, &background_jobs).await; + super::clean_up_session( + &runtime, + &v1_id, + &mut driver, + &tasks, + &background_jobs, + ) + .await; drop(binding.take()); let _ = reply.send(()); break; } None => { let v1_id = agentkit_acp::SessionId::new(session_id.to_string()); - super::clean_up_session(&v1_id, &mut driver, &tasks, &background_jobs).await; + super::clean_up_session( + &runtime, + &v1_id, + &mut driver, + &tasks, + &background_jobs, + ) + .await; break; } }, @@ -5920,6 +5934,7 @@ mod tests { Default::default(), None, loop_id.clone(), + None, ) .unwrap(); let mut driver = Agent::builder() diff --git a/src/runtime.rs b/src/runtime.rs index f08f51ce..c56c47e7 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -9,6 +9,7 @@ use std::{ time::Duration, }; +use agenthooksprotocol::SessionStartInputTrigger; use agentkit_acp::{AcpIntegration, AcpRuntimeError}; use agentkit_core::{ CancellationController, CancellationHandle, FinishReason, Item, ItemKind, MetadataMap, Part, @@ -36,6 +37,7 @@ use tracing::Instrument as _; use crate::{ acp_child::{AcpHarnesses, BUILTIN_HARNESS, ChildConfig}, + ahp::{HookedAdapter, HookedSession}, provider::{ ModelSelection, ProviderKind, ReasoningEffort, SelectableAdapter, SelectableSession, }, @@ -466,7 +468,7 @@ impl Drop for SessionClaim { pub(crate) struct AcpDriver { pub subagents: Subagents, - pub driver: LoopDriver, + pub driver: LoopDriver>, pub skills: Vec, pub tasks: TaskManagerHandle, pub background_jobs: BackgroundJobs, @@ -520,6 +522,7 @@ pub struct Runtime { skills: Arc, skill_package_roots: Vec, skill_directories: Vec, + hooks: Option>, } impl Runtime { @@ -635,6 +638,7 @@ impl Runtime { skills, skill_package_roots: Vec::new(), skill_directories: Vec::new(), + hooks: None, })) } @@ -736,6 +740,14 @@ impl Runtime { Ok(Arc::new(runtime)) } + /// Loads agent hook registrations from the global and project `.agents` directories. + pub fn with_hooks(runtime: Arc) -> Result, String> { + let mut runtime = Arc::try_unwrap(runtime) + .map_err(|_| "could not configure agent hooks after runtime was shared".to_string())?; + runtime.hooks = crate::ahp::Hooks::load(&runtime.root); + Ok(Arc::new(runtime)) + } + /// Enables evaluation only after explicit user opt-in and credential resolution. pub fn with_eval(runtime: Arc, enabled: bool) -> Result, String> { let mut runtime = Arc::try_unwrap(runtime) @@ -754,6 +766,10 @@ impl Runtime { &self.root } + pub(crate) fn hooks(&self) -> Option<&crate::ahp::Hooks> { + self.hooks.as_deref() + } + pub(crate) fn supports_terminal_authentication(&self, provider: ProviderKind) -> bool { self.credential_storage.is_persistent() && (provider != ProviderKind::OpenRouter @@ -1246,6 +1262,9 @@ impl Runtime { children.register(observe_shared(skill_tool)); } } + if let Some(hooks) = &self.hooks { + children = crate::ahp::hooked_tools(children, hooks); + } let hidden_tools = children.clone(); let compose = ComposeTool::wrap(children) .with_config( @@ -1361,6 +1380,7 @@ impl Runtime { self.agentkit_telemetry(), Some(opened.observer.clone()), format!("compaction-{}", crate::session::new_id()), + self.hooks.clone(), ) .map_err(|error| { record_runtime_failure( @@ -1383,10 +1403,19 @@ impl Runtime { error, ) })?; - let agent = Agent::builder() + let trigger = if request.resume { + SessionStartInputTrigger::Resume + } else { + SessionStartInputTrigger::Startup + }; + let mut builder = Agent::builder() .cancellation(controller.handle()) .observer(crate::fatal::RetryLog) - .model(adapter.clone()) + .model(HookedAdapter::new( + adapter.clone(), + self.hooks.clone(), + trigger, + )) .telemetry(self.agentkit_telemetry()) .add_tool_source(self.compose_with_jobs( 0, @@ -1395,7 +1424,15 @@ impl Runtime { skills, )) .task_manager(task_manager) - .mutator(compactor) + .mutator(compactor); + if let Some(hooks) = &self.hooks { + builder = builder.mutator(crate::ahp::PromptHook::new( + Arc::clone(hooks), + "cli", + crate::ahp::PromptRecord::Session(opened.observer.clone()), + )); + } + let agent = builder .transcript_observer(opened.observer) .transcript(opened.transcript) .input(vec![Item::text(ItemKind::User, prompt)]) @@ -1462,6 +1499,11 @@ impl Runtime { for task in tasks.list_running().await { let _ = tasks.cancel(task.id).await; } + if let Some(hooks) = &self.hooks { + hooks + .session_end(&session_id, crate::ahp::session_outcome(&result)) + .await; + } match result { Ok(output) => Ok(output), Err(error) => Err(record_loop_failure( @@ -1525,6 +1567,7 @@ impl Runtime { self.agentkit_telemetry(), None, format!("compaction-{session}"), + self.hooks.clone(), ) .map_err(LoopError::InvalidState)? .mirroring(recorded.clone()); @@ -1534,7 +1577,11 @@ impl Runtime { .with_transient_parent(recorded.clone()) .with_prompt_cache_key(crate::tools::subagent::prompt_cache_key(&session)); let builder = Agent::builder() - .model(self.adapter.clone()) + .model(HookedAdapter::new( + self.adapter.clone(), + self.hooks.clone(), + SessionStartInputTrigger::Startup, + )) .telemetry(self.agentkit_telemetry()) .add_tool_source(self.compose_with_jobs( depth, @@ -1543,11 +1590,20 @@ impl Runtime { skills, )) .task_manager(background_task_manager()) - .mutator(compactor) + .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()), + )), + None => builder, + }; + let builder = builder .transcript_observer(recorded) .transcript(transcript) - .input(vec![Item::text(ItemKind::User, prompt)]); - let builder = builder.cancellation(controller.handle()); + .input(vec![Item::text(ItemKind::User, prompt)]) + .cancellation(controller.handle()); let mut driver = builder .build()? .start( @@ -1557,7 +1613,13 @@ impl Runtime { ), ) .await?; - drive(&mut driver).await + let result = drive(&mut driver).await; + if let Some(hooks) = &self.hooks { + hooks + .session_end(&session, crate::ahp::session_outcome(&result)) + .await; + } + result } pub(crate) fn claim_session(self: &Arc) -> Result { @@ -1713,6 +1775,7 @@ impl Runtime { self.agentkit_telemetry(), Some(opened.observer.clone()), format!("compaction-{}", crate::session::new_id()), + self.hooks.clone(), ) .map_err(AcpRuntimeError::Loop)?; let prompt_cache_key = crate::tools::subagent::prompt_cache_key(&session_id); @@ -1731,8 +1794,19 @@ impl Runtime { if context.response_attempt_replacement { session_config = session_config.with_response_attempt_supersession(); } - let driver = Agent::builder() - .model(adapter.clone()) + let trigger = if is_fork { + SessionStartInputTrigger::Fork + } else if request.resume { + SessionStartInputTrigger::Resume + } else { + SessionStartInputTrigger::Startup + }; + let mut builder = Agent::builder() + .model(HookedAdapter::new( + adapter.clone(), + self.hooks.clone(), + trigger, + )) .telemetry(self.agentkit_telemetry()) .add_tool_source(self.compose_with_jobs_and_mcp( self.base_depth, @@ -1742,7 +1816,15 @@ impl Runtime { mcp, )) .task_manager(task_manager) - .mutator(compactor) + .mutator(compactor); + if let Some(hooks) = &self.hooks { + builder = builder.mutator(crate::ahp::PromptHook::new( + Arc::clone(hooks), + "acp", + crate::ahp::PromptRecord::Session(opened.observer.clone()), + )); + } + let driver = builder .observer(context.integration.as_ref().clone()) .observer(crate::fatal::RetryLog) .transcript_observer(opened.observer) @@ -1997,7 +2079,15 @@ fn build_skill_registry( } fn default_skill_roots(root: &Path) -> Vec { - let mut roots = vec![root.join(".agents/skills")]; + agents_directories(root) + .into_iter() + .map(|directory| directory.join("skills")) + .collect() +} + +/// Project-local `.agents`, then the user's `~/.agents` when a home is set. +pub(crate) fn agents_directories(root: &Path) -> Vec { + let mut directories = vec![root.join(".agents")]; #[cfg(unix)] let home = std::env::var_os("HOME").filter(|value| !value.is_empty()); #[cfg(not(unix))] @@ -2005,9 +2095,9 @@ fn default_skill_roots(root: &Path) -> Vec { .filter(|value| !value.is_empty()) .or_else(|| std::env::var_os("HOME").filter(|value| !value.is_empty())); if let Some(home) = home { - roots.push(PathBuf::from(home).join(".agents/skills")); + directories.push(PathBuf::from(home).join(".agents")); } - roots + directories } pub struct ComposeOnly { @@ -2856,7 +2946,9 @@ fn record_loop_failure( } } -async fn drive(driver: &mut LoopDriver) -> Result { +async fn drive( + driver: &mut LoopDriver>, +) -> Result { drive_with_tasks(driver, None).await }