Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 29 additions & 1 deletion rust/src/hooks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Comment on lines +27 to +28
/// 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.
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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 {
Expand Down
18 changes: 16 additions & 2 deletions rust/src/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand All @@ -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" => {
Expand Down
28 changes: 24 additions & 4 deletions rust/tests/session_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<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()),
Expand All @@ -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::<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
Expand Down Expand Up @@ -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");
Comment on lines +4158 to +4163
}

#[tokio::test]
Expand Down
Loading