diff --git a/rust/src/hooks.rs b/rust/src/hooks.rs index 4986d6cb18..05481f1aea 100644 --- a/rust/src/hooks.rs +++ b/rust/src/hooks.rs @@ -19,6 +19,19 @@ use crate::types::SessionId; pub struct HookContext { /// The session this hook was triggered in. pub session_id: SessionId, + /// JSON-RPC request ID for this hook invocation. + pub request_id: u64, +} + +/// Hook response that was successfully written back to the CLI. +#[derive(Debug, Clone)] +pub struct HookResponseSent { + /// The session this hook was triggered in. + pub session_id: SessionId, + /// JSON-RPC request ID for this hook invocation. + pub request_id: u64, + /// Runtime hook type, such as `userPromptSubmitted` or `preToolUse`. + pub hook_type: String, } /// Input for the `preToolUse` hook — received before a tool executes. @@ -570,6 +583,9 @@ pub trait SessionHooks: Send + Sync + 'static { } } + /// Called after a hook response is successfully written back to the CLI. + async fn on_hook_response_sent(&self, _response: HookResponseSent) {} + /// Called before a tool executes. Return `Some(output)` to approve/deny /// or modify the call, or `None` (default) to pass through unchanged. async fn on_pre_tool_use( @@ -680,14 +696,26 @@ pub trait SessionHooks: Send + Sync + 'static { /// Returns `Ok(Value)` shaped like `{ "output": ... }` on success. /// If no hook is registered ([`HookOutput::None`]), the output is an empty /// object: `{ "output": {} }`. -pub(crate) async fn dispatch_hook( +#[cfg(test)] +async fn dispatch_hook( + hooks: &dyn SessionHooks, + session_id: &SessionId, + hook_type: &str, + raw_input: Value, +) -> Result { + dispatch_hook_for_request(hooks, session_id, 0, hook_type, raw_input).await +} + +pub(crate) async fn dispatch_hook_for_request( hooks: &dyn SessionHooks, session_id: &SessionId, + request_id: u64, hook_type: &str, raw_input: Value, ) -> Result { let ctx = HookContext { session_id: session_id.clone(), + request_id, }; let event = match hook_type { diff --git a/rust/src/session.rs b/rust/src/session.rs index b9d2173055..fa27dfdce0 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -2385,7 +2385,11 @@ async fn handle_request( .unwrap_or(Value::Object(Default::default())); let rpc_result = if let Some(hooks) = hooks { - match crate::hooks::dispatch_hook(hooks, &sid, hook_type, input).await { + match crate::hooks::dispatch_hook_for_request( + hooks, &sid, request.id, hook_type, input, + ) + .await + { Ok(output) => output, Err(e) => { warn!(error = %e, hook_type = hook_type, "hook dispatch failed"); @@ -2402,7 +2406,17 @@ async fn handle_request( result: Some(rpc_result), error: None, }; - let _ = client.send_response(&rpc_response).await; + if client.send_response(&rpc_response).await.is_ok() + && let Some(hooks) = hooks + { + hooks + .on_hook_response_sent(crate::hooks::HookResponseSent { + session_id: sid, + request_id: request.id, + hook_type: hook_type.to_string(), + }) + .await; + } } "userInput.request" => { diff --git a/rust/tests/session_test.rs b/rust/tests/session_test.rs index a51d61910d..a4b3d563f8 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -4090,14 +4090,19 @@ async fn create_session_pair_with_hooks( #[tokio::test] async fn hooks_invoke_dispatches_to_session_hooks() { - use github_copilot_sdk::hooks::{HookEvent, HookOutput, PreToolUseOutput, SessionHooks}; + use github_copilot_sdk::hooks::{ + HookEvent, HookOutput, HookResponseSent, PreToolUseOutput, SessionHooks, + }; - struct PolicyHooks; + struct PolicyHooks { + response_sent: tokio::sync::mpsc::UnboundedSender, + } #[async_trait] impl SessionHooks for PolicyHooks { async fn on_hook(&self, event: HookEvent) -> HookOutput { match event { - HookEvent::PreToolUse { input, .. } => { + HookEvent::PreToolUse { input, ctx } => { + assert_eq!(ctx.request_id, 300); if input.tool_name == "rm" { HookOutput::PreToolUse(PreToolUseOutput { permission_decision: Some("deny".to_string()), @@ -4111,9 +4116,18 @@ async fn hooks_invoke_dispatches_to_session_hooks() { _ => HookOutput::None, } } + + async fn on_hook_response_sent(&self, response: HookResponseSent) { + self.response_sent.send(response).unwrap(); + } } - let (_session, mut server) = create_session_pair_with_hooks(Arc::new(PolicyHooks)).await; + let (response_sent_tx, mut response_sent_rx) = + tokio::sync::mpsc::unbounded_channel::(); + let (_session, mut server) = create_session_pair_with_hooks(Arc::new(PolicyHooks { + response_sent: response_sent_tx, + })) + .await; // Send a hooks.invoke request for a denied tool server @@ -4141,6 +4155,12 @@ async fn hooks_invoke_dispatches_to_session_hooks() { response["result"]["output"]["permissionDecisionReason"], "destructive" ); + let response_sent = timeout(TIMEOUT, response_sent_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(response_sent.request_id, 300); + assert_eq!(response_sent.hook_type, "preToolUse"); } #[tokio::test]