Notify hooks after response delivery

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: gimenete

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<Value, crate::Error> {
+    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<Value, crate::Error> {
     let ctx = HookContext {
         session_id: session_id.clone(),
+        request_id,
     };
 
     let event = match hook_type {

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" => {

rust/tests/session_test.rs

@@ -4144,14 +4144,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<HookResponseSent>,
+    }
     #[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()),
@@ -4165,9 +4170,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::<HookResponseSent>();
+    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
@@ -4195,6 +4209,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]