diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 21daadb74..d1f4305b3 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -46,6 +46,7 @@ use codex_hooks::HookEvent; use codex_hooks::HookEventAfterAgent; use codex_hooks::HookPayload; use codex_hooks::Hooks; +use codex_hooks::HooksConfig; use codex_network_proxy::NetworkProxy; use codex_protocol::ThreadId; use codex_protocol::approvals::ExecPolicyAmendment; @@ -1099,7 +1100,9 @@ impl Session { Arc::clone(&config), Arc::clone(&auth_manager), ), - hooks: Hooks::new(config.notify.clone()), + hooks: Hooks::new(HooksConfig { + legacy_notify_argv: config.notify.clone(), + }), rollout: Mutex::new(rollout_recorder), user_shell: Arc::new(default_shell), shell_snapshot_tx, @@ -6282,7 +6285,9 @@ mod tests { Arc::clone(&config), Arc::clone(&auth_manager), ), - hooks: Hooks::new(config.notify.clone()), + hooks: Hooks::new(HooksConfig { + legacy_notify_argv: config.notify.clone(), + }), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), shell_snapshot_tx: watch::channel(None).0, @@ -6425,7 +6430,9 @@ mod tests { Arc::clone(&config), Arc::clone(&auth_manager), ), - hooks: Hooks::new(config.notify.clone()), + hooks: Hooks::new(HooksConfig { + legacy_notify_argv: config.notify.clone(), + }), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), shell_snapshot_tx: watch::channel(None).0, diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 3dca006d3..d54b0f463 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use std::time::Instant; use crate::client_common::tools::ToolSpec; use crate::function_tool::FunctionCallError; @@ -10,6 +11,12 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use async_trait::async_trait; +use codex_hooks::HookEvent; +use codex_hooks::HookEventAfterToolUse; +use codex_hooks::HookPayload; +use codex_hooks::HookToolInput; +use codex_hooks::HookToolInputLocalShell; +use codex_hooks::HookToolKind; use codex_protocol::models::ResponseInputItem; use codex_utils_readiness::Readiness; use tracing::warn; @@ -121,8 +128,11 @@ impl ToolRegistry { return Err(FunctionCallError::Fatal(message)); } + let is_mutating = handler.is_mutating(&invocation).await; let output_cell = tokio::sync::Mutex::new(None); + let invocation_for_tool = invocation.clone(); + let started = Instant::now(); let result = otel .log_tool_result_with_tags( tool_name.as_ref(), @@ -132,14 +142,13 @@ impl ToolRegistry { || { let handler = handler.clone(); let output_cell = &output_cell; - let invocation = invocation; async move { - if handler.is_mutating(&invocation).await { + if is_mutating { tracing::trace!("waiting for tool gate"); - invocation.turn.tool_call_gate.wait_ready().await; + invocation_for_tool.turn.tool_call_gate.wait_ready().await; tracing::trace!("tool gate released"); } - match handler.handle(invocation).await { + match handler.handle(invocation_for_tool).await { Ok(output) => { let preview = output.log_preview(); let success = output.success_for_logging(); @@ -153,6 +162,20 @@ impl ToolRegistry { }, ) .await; + let duration = started.elapsed(); + let (output_preview, success) = match &result { + Ok((preview, success)) => (preview.clone(), *success), + Err(err) => (err.to_string(), false), + }; + dispatch_after_tool_use_hook(AfterToolUseHookDispatch { + invocation: &invocation, + output_preview, + success, + executed: true, + duration, + mutating: is_mutating, + }) + .await; match result { Ok(_) => { @@ -258,3 +281,87 @@ fn sandbox_policy_tag(policy: &SandboxPolicy) -> &'static str { SandboxPolicy::ExternalSandbox { .. } => "external-sandbox", } } + +// Hooks use a separate wire-facing input type so hook payload JSON stays stable +// and decoupled from core's internal tool runtime representation. +impl From<&ToolPayload> for HookToolInput { + fn from(payload: &ToolPayload) -> Self { + match payload { + ToolPayload::Function { arguments } => HookToolInput::Function { + arguments: arguments.clone(), + }, + ToolPayload::Custom { input } => HookToolInput::Custom { + input: input.clone(), + }, + ToolPayload::LocalShell { params } => HookToolInput::LocalShell { + params: HookToolInputLocalShell { + command: params.command.clone(), + workdir: params.workdir.clone(), + timeout_ms: params.timeout_ms, + sandbox_permissions: params.sandbox_permissions, + prefix_rule: params.prefix_rule.clone(), + justification: params.justification.clone(), + }, + }, + ToolPayload::Mcp { + server, + tool, + raw_arguments, + } => HookToolInput::Mcp { + server: server.clone(), + tool: tool.clone(), + arguments: raw_arguments.clone(), + }, + } + } +} + +fn hook_tool_kind(tool_input: &HookToolInput) -> HookToolKind { + match tool_input { + HookToolInput::Function { .. } => HookToolKind::Function, + HookToolInput::Custom { .. } => HookToolKind::Custom, + HookToolInput::LocalShell { .. } => HookToolKind::LocalShell, + HookToolInput::Mcp { .. } => HookToolKind::Mcp, + } +} + +struct AfterToolUseHookDispatch<'a> { + invocation: &'a ToolInvocation, + output_preview: String, + success: bool, + executed: bool, + duration: Duration, + mutating: bool, +} + +async fn dispatch_after_tool_use_hook(dispatch: AfterToolUseHookDispatch<'_>) { + let AfterToolUseHookDispatch { invocation, .. } = dispatch; + let session = invocation.session.as_ref(); + let turn = invocation.turn.as_ref(); + let tool_input = HookToolInput::from(&invocation.payload); + session + .hooks() + .dispatch(HookPayload { + session_id: session.conversation_id, + cwd: turn.cwd.clone(), + triggered_at: chrono::Utc::now(), + hook_event: HookEvent::AfterToolUse { + event: HookEventAfterToolUse { + turn_id: turn.sub_id.clone(), + call_id: invocation.call_id.clone(), + tool_name: invocation.tool_name.clone(), + tool_kind: hook_tool_kind(&tool_input), + tool_input, + executed: dispatch.executed, + success: dispatch.success, + duration_ms: u64::try_from(dispatch.duration.as_millis()).unwrap_or(u64::MAX), + mutating: dispatch.mutating, + sandbox: sandbox_tag(&turn.sandbox_policy, turn.windows_sandbox_level) + .to_string(), + sandbox_policy: sandbox_policy_tag(&turn.sandbox_policy).to_string(), + output_preview: dispatch.output_preview.clone(), + }, + }, + }) + .await; +} diff --git a/codex-rs/hooks/src/lib.rs b/codex-rs/hooks/src/lib.rs index c3d4383ba..aa7dd58ac 100644 --- a/codex-rs/hooks/src/lib.rs +++ b/codex-rs/hooks/src/lib.rs @@ -3,11 +3,16 @@ mod types; mod user_notification; pub use registry::Hooks; +pub use registry::HooksConfig; pub use registry::command_from_argv; pub use types::Hook; pub use types::HookEvent; pub use types::HookEventAfterAgent; +pub use types::HookEventAfterToolUse; pub use types::HookOutcome; pub use types::HookPayload; +pub use types::HookToolInput; +pub use types::HookToolInputLocalShell; +pub use types::HookToolKind; pub use user_notification::legacy_notify_json; pub use user_notification::notify_hook; diff --git a/codex-rs/hooks/src/registry.rs b/codex-rs/hooks/src/registry.rs index 0b91eea4e..9228bdddd 100644 --- a/codex-rs/hooks/src/registry.rs +++ b/codex-rs/hooks/src/registry.rs @@ -6,25 +6,42 @@ use crate::types::HookOutcome; use crate::types::HookPayload; #[derive(Default, Clone)] +pub struct HooksConfig { + pub legacy_notify_argv: Option>, +} + +#[derive(Clone)] pub struct Hooks { after_agent: Vec, + after_tool_use: Vec, +} + +impl Default for Hooks { + fn default() -> Self { + Self::new(HooksConfig::default()) + } } // Hooks are arbitrary, user-specified functions that are deterministically // executed after specific events in the Codex lifecycle. impl Hooks { - pub fn new(notify: Option>) -> Self { - let after_agent = notify + pub fn new(config: HooksConfig) -> Self { + let after_agent = config + .legacy_notify_argv .filter(|argv| !argv.is_empty() && !argv[0].is_empty()) .map(crate::notify_hook) .into_iter() .collect(); - Self { after_agent } + Self { + after_agent, + after_tool_use: Vec::new(), + } } fn hooks_for_event(&self, hook_event: &HookEvent) -> &[Hook] { match hook_event { HookEvent::AfterAgent { .. } => &self.after_agent, + HookEvent::AfterToolUse { .. } => &self.after_tool_use, } } @@ -70,6 +87,9 @@ mod tests { use super::*; use crate::types::HookEventAfterAgent; + use crate::types::HookEventAfterToolUse; + use crate::types::HookToolInput; + use crate::types::HookToolKind; const CWD: &str = "/tmp"; const INPUT_MESSAGE: &str = "hello"; @@ -106,8 +126,33 @@ mod tests { } } - fn hooks_for_after_agent(hooks: Vec) -> Hooks { - Hooks { after_agent: hooks } + fn after_tool_use_payload(label: &str) -> HookPayload { + HookPayload { + session_id: ThreadId::new(), + cwd: PathBuf::from(CWD), + triggered_at: Utc + .with_ymd_and_hms(2025, 1, 1, 0, 0, 0) + .single() + .expect("valid timestamp"), + hook_event: HookEvent::AfterToolUse { + event: HookEventAfterToolUse { + turn_id: format!("turn-{label}"), + call_id: format!("call-{label}"), + tool_name: "apply_patch".to_string(), + tool_kind: HookToolKind::Custom, + tool_input: HookToolInput::Custom { + input: "*** Begin Patch".to_string(), + }, + executed: true, + success: true, + duration_ms: 1, + mutating: true, + sandbox: "none".to_string(), + sandbox_policy: "danger-full-access".to_string(), + output_preview: "ok".to_string(), + }, + }, + } } #[test] @@ -138,17 +183,27 @@ mod tests { #[test] fn hooks_new_requires_program_name() { - assert!(Hooks::new(None).after_agent.is_empty()); - assert!(Hooks::new(Some(vec![])).after_agent.is_empty()); + assert!(Hooks::new(HooksConfig::default()).after_agent.is_empty()); assert!( - Hooks::new(Some(vec!["".to_string()])) - .after_agent - .is_empty() + Hooks::new(HooksConfig { + legacy_notify_argv: Some(vec![]), + }) + .after_agent + .is_empty() + ); + assert!( + Hooks::new(HooksConfig { + legacy_notify_argv: Some(vec!["".to_string()]), + }) + .after_agent + .is_empty() ); assert_eq!( - Hooks::new(Some(vec!["notify-send".to_string()])) - .after_agent - .len(), + Hooks::new(HooksConfig { + legacy_notify_argv: Some(vec!["notify-send".to_string()]), + }) + .after_agent + .len(), 1 ); } @@ -156,7 +211,10 @@ mod tests { #[tokio::test] async fn dispatch_executes_hook() { let calls = Arc::new(AtomicUsize::new(0)); - let hooks = hooks_for_after_agent(vec![counting_hook(&calls, HookOutcome::Continue)]); + let hooks = Hooks { + after_agent: vec![counting_hook(&calls, HookOutcome::Continue)], + ..Hooks::default() + }; hooks.dispatch(hook_payload("1")).await; assert_eq!(calls.load(Ordering::SeqCst), 1); @@ -172,10 +230,13 @@ mod tests { #[tokio::test] async fn dispatch_executes_multiple_hooks_for_same_event() { let calls = Arc::new(AtomicUsize::new(0)); - let hooks = hooks_for_after_agent(vec![ - counting_hook(&calls, HookOutcome::Continue), - counting_hook(&calls, HookOutcome::Continue), - ]); + let hooks = Hooks { + after_agent: vec![ + counting_hook(&calls, HookOutcome::Continue), + counting_hook(&calls, HookOutcome::Continue), + ], + ..Hooks::default() + }; hooks.dispatch(hook_payload("2")).await; assert_eq!(calls.load(Ordering::SeqCst), 2); @@ -184,15 +245,30 @@ mod tests { #[tokio::test] async fn dispatch_stops_when_hook_returns_stop() { let calls = Arc::new(AtomicUsize::new(0)); - let hooks = hooks_for_after_agent(vec![ - counting_hook(&calls, HookOutcome::Stop), - counting_hook(&calls, HookOutcome::Continue), - ]); + let hooks = Hooks { + after_agent: vec![ + counting_hook(&calls, HookOutcome::Stop), + counting_hook(&calls, HookOutcome::Continue), + ], + ..Hooks::default() + }; hooks.dispatch(hook_payload("3")).await; assert_eq!(calls.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn dispatch_executes_after_tool_use_hooks() { + let calls = Arc::new(AtomicUsize::new(0)); + let hooks = Hooks { + after_tool_use: vec![counting_hook(&calls, HookOutcome::Continue)], + ..Hooks::default() + }; + + hooks.dispatch(after_tool_use_payload("p")).await; + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + #[cfg(not(windows))] #[tokio::test] async fn hook_executes_program_with_payload_argument_unix() -> Result<()> { @@ -222,7 +298,10 @@ mod tests { let payload = hook_payload("4"); let expected = to_string(&payload)?; - let hooks = hooks_for_after_agent(vec![hook]); + let hooks = Hooks { + after_agent: vec![hook], + ..Hooks::default() + }; hooks.dispatch(payload).await; let contents = timeout(Duration::from_secs(2), async { @@ -277,7 +356,10 @@ mod tests { let payload = hook_payload("4"); let expected = to_string(&payload)?; - let hooks = hooks_for_after_agent(vec![hook]); + let hooks = Hooks { + after_agent: vec![hook], + ..Hooks::default() + }; hooks.dispatch(payload).await; let contents = timeout(Duration::from_secs(2), async { diff --git a/codex-rs/hooks/src/types.rs b/codex-rs/hooks/src/types.rs index 8cb39b5ea..a1588b897 100644 --- a/codex-rs/hooks/src/types.rs +++ b/codex-rs/hooks/src/types.rs @@ -5,6 +5,7 @@ use chrono::DateTime; use chrono::SecondsFormat; use chrono::Utc; use codex_protocol::ThreadId; +use codex_protocol::models::SandboxPermissions; use futures::future::BoxFuture; use serde::Serialize; use serde::Serializer; @@ -49,6 +50,62 @@ pub struct HookEventAfterAgent { pub last_assistant_message: Option, } +#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum HookToolKind { + Function, + Custom, + LocalShell, + Mcp, +} + +#[derive(Debug, Clone, Serialize, PartialEq)] +#[serde(rename_all = "snake_case")] +pub struct HookToolInputLocalShell { + pub command: Vec, + pub workdir: Option, + pub timeout_ms: Option, + pub sandbox_permissions: Option, + pub prefix_rule: Option>, + pub justification: Option, +} + +#[derive(Debug, Clone, Serialize, PartialEq)] +#[serde(tag = "input_type", rename_all = "snake_case")] +pub enum HookToolInput { + Function { + arguments: String, + }, + Custom { + input: String, + }, + LocalShell { + params: HookToolInputLocalShell, + }, + Mcp { + server: String, + tool: String, + arguments: String, + }, +} + +#[derive(Debug, Clone, Serialize, PartialEq)] +#[serde(rename_all = "snake_case")] +pub struct HookEventAfterToolUse { + pub turn_id: String, + pub call_id: String, + pub tool_name: String, + pub tool_kind: HookToolKind, + pub tool_input: HookToolInput, + pub executed: bool, + pub success: bool, + pub duration_ms: u64, + pub mutating: bool, + pub sandbox: String, + pub sandbox_policy: String, + pub output_preview: String, +} + fn serialize_triggered_at(value: &DateTime, serializer: S) -> Result where S: Serializer, @@ -63,6 +120,10 @@ pub enum HookEvent { #[serde(flatten)] event: HookEventAfterAgent, }, + AfterToolUse { + #[serde(flatten)] + event: HookEventAfterToolUse, + }, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -79,12 +140,17 @@ mod tests { use chrono::TimeZone; use chrono::Utc; use codex_protocol::ThreadId; + use codex_protocol::models::SandboxPermissions; use pretty_assertions::assert_eq; use serde_json::json; use super::HookEvent; use super::HookEventAfterAgent; + use super::HookEventAfterToolUse; use super::HookPayload; + use super::HookToolInput; + use super::HookToolInputLocalShell; + use super::HookToolKind; #[test] fn hook_payload_serializes_stable_wire_shape() { @@ -123,4 +189,76 @@ mod tests { assert_eq!(actual, expected); } + + #[test] + fn after_tool_use_payload_serializes_stable_wire_shape() { + let session_id = ThreadId::new(); + let payload = HookPayload { + session_id, + cwd: PathBuf::from("tmp"), + triggered_at: Utc + .with_ymd_and_hms(2025, 1, 1, 0, 0, 0) + .single() + .expect("valid timestamp"), + hook_event: HookEvent::AfterToolUse { + event: HookEventAfterToolUse { + turn_id: "turn-2".to_string(), + call_id: "call-1".to_string(), + tool_name: "local_shell".to_string(), + tool_kind: HookToolKind::LocalShell, + tool_input: HookToolInput::LocalShell { + params: HookToolInputLocalShell { + command: vec!["cargo".to_string(), "fmt".to_string()], + workdir: Some("codex-rs".to_string()), + timeout_ms: Some(60_000), + sandbox_permissions: Some(SandboxPermissions::UseDefault), + justification: None, + prefix_rule: None, + }, + }, + executed: true, + success: true, + duration_ms: 42, + mutating: true, + sandbox: "none".to_string(), + sandbox_policy: "danger-full-access".to_string(), + output_preview: "ok".to_string(), + }, + }, + }; + + let actual = serde_json::to_value(payload).expect("serialize hook payload"); + let expected = json!({ + "session_id": session_id.to_string(), + "cwd": "tmp", + "triggered_at": "2025-01-01T00:00:00Z", + "hook_event": { + "event_type": "after_tool_use", + "turn_id": "turn-2", + "call_id": "call-1", + "tool_name": "local_shell", + "tool_kind": "local_shell", + "tool_input": { + "input_type": "local_shell", + "params": { + "command": ["cargo", "fmt"], + "workdir": "codex-rs", + "timeout_ms": 60000, + "sandbox_permissions": "use_default", + "justification": null, + "prefix_rule": null, + }, + }, + "executed": true, + "success": true, + "duration_ms": 42, + "mutating": true, + "sandbox": "none", + "sandbox_policy": "danger-full-access", + "output_preview": "ok", + }, + }); + + assert_eq!(actual, expected); + } } diff --git a/codex-rs/hooks/src/user_notification.rs b/codex-rs/hooks/src/user_notification.rs index 1e33404b5..d3104a1ec 100644 --- a/codex-rs/hooks/src/user_notification.rs +++ b/codex-rs/hooks/src/user_notification.rs @@ -29,15 +29,20 @@ enum UserNotification { } pub fn legacy_notify_json(hook_event: &HookEvent, cwd: &Path) -> Result { - serde_json::to_string(&match hook_event { - HookEvent::AfterAgent { event } => UserNotification::AgentTurnComplete { - thread_id: event.thread_id.to_string(), - turn_id: event.turn_id.clone(), - cwd: cwd.display().to_string(), - input_messages: event.input_messages.clone(), - last_assistant_message: event.last_assistant_message.clone(), - }, - }) + match hook_event { + HookEvent::AfterAgent { event } => { + serde_json::to_string(&UserNotification::AgentTurnComplete { + thread_id: event.thread_id.to_string(), + turn_id: event.turn_id.clone(), + cwd: cwd.display().to_string(), + input_messages: event.input_messages.clone(), + last_assistant_message: event.last_assistant_message.clone(), + }) + } + _ => Err(serde_json::Error::io(std::io::Error::other( + "legacy notify payload is only supported for after_agent", + ))), + } } pub fn notify_hook(argv: Vec) -> Hook {