mirror of
https://github.com/openai/codex.git
synced 2026-11-12 21:20:34 +00:00
Add detached async command hooks
This commit is contained in:
@@ -0,0 +1,204 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::MutexGuard;
|
||||
|
||||
use codex_protocol::protocol::HookEventName;
|
||||
use codex_utils_output_truncation::TruncationPolicy;
|
||||
use codex_utils_output_truncation::approx_token_count;
|
||||
use codex_utils_output_truncation::formatted_truncate_text;
|
||||
use codex_utils_output_truncation::truncate_text;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tokio_util::task::TaskTracker;
|
||||
|
||||
use super::command_runner::CommandRunResult;
|
||||
use super::output_parser;
|
||||
|
||||
const ASYNC_HOOK_COMPLETION_TOKEN_LIMIT: usize = 500;
|
||||
const ASYNC_HOOK_COMPLETION_TRUNCATION_TOKEN_LIMIT: usize = 450;
|
||||
const ASYNC_HOOK_FLUSH_TOKEN_LIMIT: usize = 1_000;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct AsyncHookCompletion {
|
||||
event_name: HookEventName,
|
||||
text: String,
|
||||
}
|
||||
|
||||
/// Session-scoped FIFO for informational output from detached command hooks.
|
||||
#[derive(Clone)]
|
||||
pub struct AsyncHookOutputQueue {
|
||||
state: Arc<AsyncHookOutputState>,
|
||||
}
|
||||
|
||||
struct AsyncHookOutputState {
|
||||
pending: Mutex<VecDeque<AsyncHookCompletion>>,
|
||||
cancellation: CancellationToken,
|
||||
tasks: TaskTracker,
|
||||
}
|
||||
|
||||
/// A bounded prefix prepared for one accepted user turn.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AsyncHookOutputBatch {
|
||||
completion_count: usize,
|
||||
text: String,
|
||||
}
|
||||
|
||||
impl AsyncHookOutputBatch {
|
||||
/// Returns the merged developer context for this batch.
|
||||
pub fn into_text(self) -> String {
|
||||
self.text
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncHookOutputQueue {
|
||||
/// Records one completed hook's bounded informational output.
|
||||
pub fn push(&self, event_name: HookEventName, text: String) {
|
||||
if text.trim().is_empty() {
|
||||
return;
|
||||
}
|
||||
let text = formatted_truncate_text(
|
||||
&text,
|
||||
TruncationPolicy::Tokens(ASYNC_HOOK_COMPLETION_TRUNCATION_TOKEN_LIMIT),
|
||||
);
|
||||
let text = if approx_token_count(&text) > ASYNC_HOOK_COMPLETION_TOKEN_LIMIT {
|
||||
truncate_text(
|
||||
&text,
|
||||
TruncationPolicy::Tokens(ASYNC_HOOK_COMPLETION_TOKEN_LIMIT),
|
||||
)
|
||||
} else {
|
||||
text
|
||||
};
|
||||
self.lock_pending()
|
||||
.push_back(AsyncHookCompletion { event_name, text });
|
||||
}
|
||||
|
||||
/// Prepares a bounded FIFO prefix without removing it from the queue.
|
||||
pub fn pending_batch(&self) -> Option<AsyncHookOutputBatch> {
|
||||
let pending = self.lock_pending();
|
||||
let mut selected = Vec::new();
|
||||
for completion in pending.iter() {
|
||||
let mut candidate = selected.clone();
|
||||
candidate.push(completion.clone());
|
||||
let text = render_batch(&candidate);
|
||||
if approx_token_count(&text) > ASYNC_HOOK_FLUSH_TOKEN_LIMIT {
|
||||
break;
|
||||
}
|
||||
selected = candidate;
|
||||
}
|
||||
(!selected.is_empty()).then(|| AsyncHookOutputBatch {
|
||||
completion_count: selected.len(),
|
||||
text: render_batch(&selected),
|
||||
})
|
||||
}
|
||||
|
||||
/// Removes a previously prepared prefix after its user input is accepted.
|
||||
pub fn commit(&self, batch: &AsyncHookOutputBatch) {
|
||||
self.lock_pending().drain(..batch.completion_count);
|
||||
}
|
||||
|
||||
/// Cancels detached hook commands and waits for their tasks to stop.
|
||||
pub async fn shutdown(&self) {
|
||||
self.state.cancellation.cancel();
|
||||
self.state.tasks.close();
|
||||
self.state.tasks.wait().await;
|
||||
}
|
||||
|
||||
pub(crate) fn cancellation_token(&self) -> CancellationToken {
|
||||
self.state.cancellation.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn spawn(&self, future: impl Future<Output = ()> + Send + 'static) {
|
||||
self.state.tasks.spawn(future);
|
||||
}
|
||||
|
||||
fn lock_pending(&self) -> MutexGuard<'_, VecDeque<AsyncHookCompletion>> {
|
||||
self.state
|
||||
.pending
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AsyncHookOutputQueue {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
state: Arc::new(AsyncHookOutputState {
|
||||
pending: Mutex::new(VecDeque::new()),
|
||||
cancellation: CancellationToken::new(),
|
||||
tasks: TaskTracker::new(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn render_batch(completions: &[AsyncHookCompletion]) -> String {
|
||||
let outputs = completions
|
||||
.iter()
|
||||
.map(|completion| {
|
||||
format!(
|
||||
"<async_hook_output event=\"{:?}\">\n{}\n</async_hook_output>",
|
||||
completion.event_name, completion.text
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
format!("<async_hook_outputs>\n{outputs}\n</async_hook_outputs>")
|
||||
}
|
||||
|
||||
pub(crate) fn deliverable_output(
|
||||
event_name: HookEventName,
|
||||
run_result: &CommandRunResult,
|
||||
) -> Option<String> {
|
||||
if let Some(error) = run_result.error.as_deref() {
|
||||
return Some(format!("Async hook failed to run: {error}"));
|
||||
}
|
||||
|
||||
match run_result.exit_code {
|
||||
Some(0) => parse_stdout(event_name, &run_result.stdout),
|
||||
Some(exit_code) => Some(format!("Async hook exited with code {exit_code}")),
|
||||
None => Some("Async hook process terminated without an exit code".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_stdout(event_name: HookEventName, stdout: &str) -> Option<String> {
|
||||
let trimmed = stdout.trim();
|
||||
if trimmed.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let additional_context = match event_name {
|
||||
HookEventName::SessionStart => {
|
||||
output_parser::parse_session_start(trimmed).map(|output| output.additional_context)
|
||||
}
|
||||
HookEventName::SubagentStart => {
|
||||
output_parser::parse_subagent_start(trimmed).map(|output| output.additional_context)
|
||||
}
|
||||
HookEventName::PreToolUse => {
|
||||
output_parser::parse_pre_tool_use(trimmed).map(|output| output.additional_context)
|
||||
}
|
||||
HookEventName::PostToolUse => {
|
||||
output_parser::parse_post_tool_use(trimmed).map(|output| output.additional_context)
|
||||
}
|
||||
HookEventName::UserPromptSubmit => {
|
||||
output_parser::parse_user_prompt_submit(trimmed).map(|output| output.additional_context)
|
||||
}
|
||||
HookEventName::PermissionRequest => {
|
||||
output_parser::parse_permission_request(trimmed).map(|_| None)
|
||||
}
|
||||
HookEventName::PreCompact => output_parser::parse_pre_compact(trimmed).map(|_| None),
|
||||
HookEventName::PostCompact => output_parser::parse_post_compact(trimmed).map(|_| None),
|
||||
HookEventName::SubagentStop => output_parser::parse_subagent_stop(trimmed).map(|_| None),
|
||||
HookEventName::Stop => output_parser::parse_stop(trimmed).map(|_| None),
|
||||
};
|
||||
match additional_context {
|
||||
Some(additional_context) => additional_context.filter(|context| !context.trim().is_empty()),
|
||||
None => Some(format!(
|
||||
"Async {event_name:?} hook returned invalid JSON output"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "async_output_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,326 @@
|
||||
use std::fs;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::Duration;
|
||||
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::protocol::HookEventName;
|
||||
use codex_protocol::protocol::HookSource;
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use codex_utils_output_truncation::approx_token_count;
|
||||
use pretty_assertions::assert_eq;
|
||||
use tempfile::tempdir;
|
||||
use tokio::time::sleep;
|
||||
use tokio::time::timeout;
|
||||
|
||||
use super::ASYNC_HOOK_COMPLETION_TOKEN_LIMIT;
|
||||
use super::ASYNC_HOOK_FLUSH_TOKEN_LIMIT;
|
||||
use super::AsyncHookOutputQueue;
|
||||
use super::deliverable_output;
|
||||
use crate::engine::CommandShell;
|
||||
use crate::engine::ConfiguredHandler;
|
||||
use crate::engine::command_runner::CommandRunResult;
|
||||
use crate::events::permission_request;
|
||||
use crate::events::post_tool_use;
|
||||
use crate::events::pre_tool_use;
|
||||
|
||||
#[test]
|
||||
fn async_output_delivers_only_event_context() {
|
||||
let pre = output(
|
||||
HookEventName::PreToolUse,
|
||||
r#"{"continue":false,"systemMessage":"ignore","decision":"block","reason":"ignore","hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"deny","updatedInput":{"command":"rewrite"},"additionalContext":"pre context"}}"#,
|
||||
);
|
||||
let permission = output(
|
||||
HookEventName::PermissionRequest,
|
||||
r#"{"continue":false,"hookSpecificOutput":{"hookEventName":"PermissionRequest","decision":{"behavior":"deny","message":"ignore"}}}"#,
|
||||
);
|
||||
let post = output(
|
||||
HookEventName::PostToolUse,
|
||||
r#"{"continue":false,"decision":"block","reason":"ignore","hookSpecificOutput":{"hookEventName":"PostToolUse","additionalContext":"post context"}}"#,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
(pre.as_deref(), permission, post.as_deref()),
|
||||
(Some("pre context"), None, Some("post context"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn async_output_surfaces_parse_and_runtime_failures() {
|
||||
let invalid = deliverable_output(
|
||||
HookEventName::PreToolUse,
|
||||
&result(Some(0), "plain stdout", "", /*error*/ None),
|
||||
);
|
||||
let nonzero = deliverable_output(
|
||||
HookEventName::PreToolUse,
|
||||
&result(Some(2), "", "denied", /*error*/ None),
|
||||
);
|
||||
let runtime = deliverable_output(
|
||||
HookEventName::PreToolUse,
|
||||
&result(
|
||||
/*exit_code*/ None,
|
||||
"",
|
||||
"",
|
||||
Some("spawn failed".to_string()),
|
||||
),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
(invalid, nonzero, runtime),
|
||||
(
|
||||
Some("Async PreToolUse hook returned invalid JSON output".to_string()),
|
||||
Some("Async hook exited with code 2".to_string()),
|
||||
Some("Async hook failed to run: spawn failed".to_string()),
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
#[tokio::test]
|
||||
async fn async_tool_hooks_do_not_return_control_or_lifecycle() {
|
||||
let temp = tempdir().expect("create temp dir");
|
||||
let script_path = temp.path().join("async_controls.py");
|
||||
fs::write(
|
||||
&script_path,
|
||||
r#"import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
name = sys.argv[1]
|
||||
outputs = {
|
||||
"pre_block": {"hookSpecificOutput": {"hookEventName": "PreToolUse", "permissionDecision": "deny", "permissionDecisionReason": "blocked", "additionalContext": "pre block context"}},
|
||||
"pre_rewrite": {"hookSpecificOutput": {"hookEventName": "PreToolUse", "permissionDecision": "allow", "updatedInput": {"command": "rewrite"}, "additionalContext": "pre rewrite context"}},
|
||||
"permission": {"hookSpecificOutput": {"hookEventName": "PermissionRequest", "decision": {"behavior": "deny", "message": "denied"}}},
|
||||
"post": {"continue": False, "decision": "block", "reason": "stop", "hookSpecificOutput": {"hookEventName": "PostToolUse", "additionalContext": "post context"}},
|
||||
}
|
||||
json.load(sys.stdin)
|
||||
print(json.dumps(outputs[name]))
|
||||
(Path(__file__).parent / f"{name}.completed").write_text("done", encoding="utf-8")
|
||||
"#,
|
||||
)
|
||||
.expect("write async controls script");
|
||||
let cwd = AbsolutePathBuf::try_from(temp.path().to_path_buf()).expect("absolute temp dir");
|
||||
let queue = AsyncHookOutputQueue::default();
|
||||
let shell = CommandShell {
|
||||
program: String::new(),
|
||||
args: Vec::new(),
|
||||
};
|
||||
let pre_request = pre_tool_use::PreToolUseRequest {
|
||||
session_id: ThreadId::new(),
|
||||
turn_id: "turn-1".to_string(),
|
||||
subagent: None,
|
||||
cwd: cwd.clone(),
|
||||
transcript_path: None,
|
||||
model: "gpt-test".to_string(),
|
||||
permission_mode: "default".to_string(),
|
||||
tool_name: "Bash".to_string(),
|
||||
matcher_aliases: Vec::new(),
|
||||
tool_use_id: "tool-1".to_string(),
|
||||
tool_input: serde_json::json!({"command": "original"}),
|
||||
};
|
||||
let pre_handlers = [
|
||||
async_handler(HookEventName::PreToolUse, &script_path, "pre_block"),
|
||||
async_handler(HookEventName::PreToolUse, &script_path, "pre_rewrite"),
|
||||
];
|
||||
assert!(pre_tool_use::preview(&pre_handlers, &pre_request).is_empty());
|
||||
let pre = pre_tool_use::run(&pre_handlers, &shell, &queue, pre_request).await;
|
||||
assert_eq!(
|
||||
(
|
||||
pre.hook_events,
|
||||
pre.should_block,
|
||||
pre.block_reason,
|
||||
pre.additional_contexts,
|
||||
pre.updated_input,
|
||||
),
|
||||
(Vec::new(), false, None, Vec::new(), None)
|
||||
);
|
||||
|
||||
let permission_request = permission_request::PermissionRequestRequest {
|
||||
session_id: ThreadId::new(),
|
||||
turn_id: "turn-1".to_string(),
|
||||
subagent: None,
|
||||
cwd: temp.path().to_path_buf(),
|
||||
transcript_path: None,
|
||||
model: "gpt-test".to_string(),
|
||||
permission_mode: "default".to_string(),
|
||||
tool_name: "Bash".to_string(),
|
||||
matcher_aliases: Vec::new(),
|
||||
run_id_suffix: "tool-1".to_string(),
|
||||
tool_input: serde_json::json!({"command": "original"}),
|
||||
};
|
||||
let permission_handlers = [async_handler(
|
||||
HookEventName::PermissionRequest,
|
||||
&script_path,
|
||||
"permission",
|
||||
)];
|
||||
assert!(permission_request::preview(&permission_handlers, &permission_request).is_empty());
|
||||
let permission =
|
||||
permission_request::run(&permission_handlers, &shell, &queue, permission_request).await;
|
||||
assert_eq!(
|
||||
(permission.hook_events, permission.decision),
|
||||
(Vec::new(), None)
|
||||
);
|
||||
|
||||
let post_request = post_tool_use::PostToolUseRequest {
|
||||
session_id: ThreadId::new(),
|
||||
turn_id: "turn-1".to_string(),
|
||||
subagent: None,
|
||||
cwd,
|
||||
transcript_path: None,
|
||||
model: "gpt-test".to_string(),
|
||||
permission_mode: "default".to_string(),
|
||||
tool_name: "Bash".to_string(),
|
||||
matcher_aliases: Vec::new(),
|
||||
tool_use_id: "tool-1".to_string(),
|
||||
tool_input: serde_json::json!({"command": "original"}),
|
||||
tool_response: serde_json::json!({"ok": true}),
|
||||
};
|
||||
let post_handlers = [async_handler(
|
||||
HookEventName::PostToolUse,
|
||||
&script_path,
|
||||
"post",
|
||||
)];
|
||||
assert!(post_tool_use::preview(&post_handlers, &post_request).is_empty());
|
||||
let post = post_tool_use::run(&post_handlers, &shell, &queue, post_request).await;
|
||||
assert_eq!(
|
||||
(
|
||||
post.hook_events,
|
||||
post.should_stop,
|
||||
post.stop_reason,
|
||||
post.additional_contexts,
|
||||
post.feedback_message,
|
||||
),
|
||||
(Vec::new(), false, None, Vec::new(), None)
|
||||
);
|
||||
|
||||
timeout(Duration::from_secs(10), async {
|
||||
for name in ["pre_block", "pre_rewrite", "permission", "post"] {
|
||||
while !temp.path().join(format!("{name}.completed")).exists() {
|
||||
sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("async control hooks complete");
|
||||
let delivered = queue
|
||||
.pending_batch()
|
||||
.expect("informational async outputs")
|
||||
.into_text();
|
||||
assert!(delivered.contains("pre block context"));
|
||||
assert!(delivered.contains("pre rewrite context"));
|
||||
assert!(delivered.contains("post context"));
|
||||
assert!(!delivered.contains("denied"));
|
||||
queue.shutdown().await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_preserves_duplicate_completions_until_commit() {
|
||||
let queue = AsyncHookOutputQueue::default();
|
||||
queue.push(HookEventName::PreToolUse, "same".to_string());
|
||||
queue.push(HookEventName::PostToolUse, "same".to_string());
|
||||
queue.push(HookEventName::Stop, "last".to_string());
|
||||
|
||||
let batch = queue.pending_batch().expect("queued output batch");
|
||||
assert_eq!(queue.pending_batch(), Some(batch.clone()));
|
||||
let text = batch.clone().into_text();
|
||||
assert_eq!(text.matches("same").count(), 2);
|
||||
assert!(
|
||||
["PreToolUse", "PostToolUse", "last"]
|
||||
.map(|needle| text.find(needle).expect("queued output"))
|
||||
.is_sorted()
|
||||
);
|
||||
queue.commit(&batch);
|
||||
assert!(queue.pending_batch().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_bounds_items_and_flushes_a_contiguous_prefix() {
|
||||
let queue = AsyncHookOutputQueue::default();
|
||||
let large = "large-output ".repeat(ASYNC_HOOK_COMPLETION_TOKEN_LIMIT);
|
||||
for event_name in [
|
||||
HookEventName::PreToolUse,
|
||||
HookEventName::PostToolUse,
|
||||
HookEventName::Stop,
|
||||
] {
|
||||
queue.push(event_name, large.clone());
|
||||
}
|
||||
queue.push(HookEventName::SessionStart, "small-tail".to_string());
|
||||
assert!(queue.lock_pending().iter().all(|completion| {
|
||||
approx_token_count(&completion.text) <= ASYNC_HOOK_COMPLETION_TOKEN_LIMIT
|
||||
}));
|
||||
|
||||
let mut completed = 0;
|
||||
loop {
|
||||
let batch = queue.pending_batch().expect("bounded output batch");
|
||||
let text = batch.clone().into_text();
|
||||
assert!(approx_token_count(&text) <= ASYNC_HOOK_FLUSH_TOKEN_LIMIT);
|
||||
completed += batch.completion_count;
|
||||
queue.commit(&batch);
|
||||
if let Some(tail) = text.find("small-tail") {
|
||||
if let Some(large) = text.find("tokens truncated") {
|
||||
assert!(large < tail);
|
||||
}
|
||||
assert_eq!(completed, 4);
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(queue.pending_batch().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_cancels_and_joins_detached_tasks() {
|
||||
let queue = AsyncHookOutputQueue::default();
|
||||
let cancellation = queue.cancellation_token();
|
||||
let stopped = Arc::new(AtomicBool::new(false));
|
||||
let task_stopped = Arc::clone(&stopped);
|
||||
queue.spawn(async move {
|
||||
cancellation.cancelled().await;
|
||||
task_stopped.store(true, Ordering::SeqCst);
|
||||
});
|
||||
|
||||
queue.shutdown().await;
|
||||
|
||||
assert!(stopped.load(Ordering::SeqCst));
|
||||
}
|
||||
|
||||
fn output(event_name: HookEventName, stdout: &str) -> Option<String> {
|
||||
deliverable_output(event_name, &result(Some(0), stdout, "", /*error*/ None))
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
fn async_handler(
|
||||
event_name: HookEventName,
|
||||
script_path: &std::path::Path,
|
||||
name: &str,
|
||||
) -> ConfiguredHandler {
|
||||
ConfiguredHandler {
|
||||
event_name,
|
||||
matcher: Some("^Bash$".to_string()),
|
||||
command: format!("python3 {} {name}", script_path.display()),
|
||||
timeout_sec: 5,
|
||||
r#async: true,
|
||||
status_message: None,
|
||||
source_path: AbsolutePathBuf::try_from(script_path.to_path_buf()).expect("absolute script"),
|
||||
source: HookSource::User,
|
||||
display_order: 0,
|
||||
env: std::collections::HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn result(
|
||||
exit_code: Option<i32>,
|
||||
stdout: &str,
|
||||
stderr: &str,
|
||||
error: Option<String>,
|
||||
) -> CommandRunResult {
|
||||
CommandRunResult {
|
||||
started_at: 1,
|
||||
completed_at: 2,
|
||||
duration_ms: 1,
|
||||
exit_code,
|
||||
stdout: stdout.to_string(),
|
||||
stderr: stderr.to_string(),
|
||||
error,
|
||||
}
|
||||
}
|
||||
@@ -465,13 +465,6 @@ fn append_matcher_groups(
|
||||
} else {
|
||||
command
|
||||
};
|
||||
if r#async {
|
||||
warnings.push(format!(
|
||||
"skipping async hook in {}: async hooks are not supported yet",
|
||||
source.path.display()
|
||||
));
|
||||
continue;
|
||||
}
|
||||
if command.trim().is_empty() {
|
||||
warnings.push(format!(
|
||||
"skipping empty hook command in {}",
|
||||
@@ -529,6 +522,7 @@ fn append_matcher_groups(
|
||||
matcher: matcher.map(ToOwned::to_owned),
|
||||
command,
|
||||
timeout_sec,
|
||||
r#async,
|
||||
status_message,
|
||||
source_path: source.path.clone(),
|
||||
source: source.source,
|
||||
@@ -780,6 +774,7 @@ mod tests {
|
||||
matcher: None,
|
||||
command: "echo hello".to_string(),
|
||||
timeout_sec: 600,
|
||||
r#async: false,
|
||||
status_message: None,
|
||||
source_path: source_path.clone(),
|
||||
source: hook_source(),
|
||||
@@ -790,8 +785,9 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pre_tool_use_keeps_valid_matcher_during_discovery() {
|
||||
fn async_pre_tool_use_is_discovered_without_warning() {
|
||||
let mut handlers = Vec::new();
|
||||
let mut hook_entries = Vec::new();
|
||||
let mut warnings = Vec::new();
|
||||
let mut display_order = 0;
|
||||
let source_path = source_path();
|
||||
@@ -799,12 +795,21 @@ mod tests {
|
||||
|
||||
append_matcher_groups(
|
||||
&mut handlers,
|
||||
&mut Vec::new(),
|
||||
&mut hook_entries,
|
||||
&mut warnings,
|
||||
&mut display_order,
|
||||
&hook_handler_source(&source_path, &hook_states),
|
||||
HookEventName::PreToolUse,
|
||||
vec![command_group(Some("^Bash$"))],
|
||||
vec![MatcherGroup {
|
||||
matcher: Some("^Bash$".to_string()),
|
||||
hooks: vec![HookHandlerConfig::Command {
|
||||
command: "echo hello".to_string(),
|
||||
command_windows: None,
|
||||
timeout_sec: None,
|
||||
r#async: true,
|
||||
status_message: None,
|
||||
}],
|
||||
}],
|
||||
);
|
||||
|
||||
assert_eq!(warnings, Vec::<String>::new());
|
||||
@@ -815,6 +820,7 @@ mod tests {
|
||||
matcher: Some("^Bash$".to_string()),
|
||||
command: "echo hello".to_string(),
|
||||
timeout_sec: 600,
|
||||
r#async: true,
|
||||
status_message: None,
|
||||
source_path: source_path.clone(),
|
||||
source: hook_source(),
|
||||
@@ -822,6 +828,7 @@ mod tests {
|
||||
env: std::collections::HashMap::new(),
|
||||
}]
|
||||
);
|
||||
assert_eq!(hook_entries.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use futures::StreamExt;
|
||||
use futures::stream::FuturesUnordered;
|
||||
@@ -13,6 +14,8 @@ use codex_protocol::protocol::HookScope;
|
||||
|
||||
use super::CommandShell;
|
||||
use super::ConfiguredHandler;
|
||||
use super::async_output::AsyncHookOutputQueue;
|
||||
use super::async_output::deliverable_output;
|
||||
use super::command_runner::CommandRunResult;
|
||||
use super::command_runner::run_command;
|
||||
use crate::events::common::matches_matcher;
|
||||
@@ -33,6 +36,17 @@ pub(crate) fn select_handlers(
|
||||
select_handlers_for_matcher_inputs(handlers, event_name, &matcher_inputs)
|
||||
}
|
||||
|
||||
pub(crate) fn select_sync_handlers(
|
||||
handlers: &[ConfiguredHandler],
|
||||
event_name: HookEventName,
|
||||
matcher_input: Option<&str>,
|
||||
) -> Vec<ConfiguredHandler> {
|
||||
select_handlers(handlers, event_name, matcher_input)
|
||||
.into_iter()
|
||||
.filter(|handler| !handler.r#async)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn select_handlers_for_matcher_inputs(
|
||||
handlers: &[ConfiguredHandler],
|
||||
event_name: HookEventName,
|
||||
@@ -67,6 +81,80 @@ pub(crate) fn select_handlers_for_matcher_inputs(
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn select_sync_handlers_for_matcher_inputs(
|
||||
handlers: &[ConfiguredHandler],
|
||||
event_name: HookEventName,
|
||||
matcher_inputs: &[&str],
|
||||
) -> Vec<ConfiguredHandler> {
|
||||
select_handlers_for_matcher_inputs(handlers, event_name, matcher_inputs)
|
||||
.into_iter()
|
||||
.filter(|handler| !handler.r#async)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn partition_handlers(
|
||||
handlers: Vec<ConfiguredHandler>,
|
||||
) -> (Vec<ConfiguredHandler>, Vec<ConfiguredHandler>) {
|
||||
handlers.into_iter().partition(|handler| !handler.r#async)
|
||||
}
|
||||
|
||||
fn spawn_async_handlers(
|
||||
shell: &CommandShell,
|
||||
output_queue: &AsyncHookOutputQueue,
|
||||
handlers: Vec<ConfiguredHandler>,
|
||||
input_json: String,
|
||||
cwd: &Path,
|
||||
) {
|
||||
for handler in handlers {
|
||||
let shell = shell.clone();
|
||||
let output_queue = output_queue.clone();
|
||||
let cancellation = output_queue.cancellation_token();
|
||||
let input_json = input_json.clone();
|
||||
let cwd = PathBuf::from(cwd);
|
||||
output_queue.clone().spawn(async move {
|
||||
tokio::select! {
|
||||
_ = cancellation.cancelled() => {}
|
||||
run_result = run_command(&shell, &handler, &input_json, &cwd) => {
|
||||
if let Some(output) = deliverable_output(handler.event_name, &run_result) {
|
||||
output_queue.push(handler.event_name, output);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_async_failures(
|
||||
output_queue: &AsyncHookOutputQueue,
|
||||
handlers: Vec<ConfiguredHandler>,
|
||||
error: String,
|
||||
) {
|
||||
for handler in handlers {
|
||||
let output_queue = output_queue.clone();
|
||||
let cancellation = output_queue.cancellation_token();
|
||||
let error = error.clone();
|
||||
output_queue.clone().spawn(async move {
|
||||
if !cancellation.is_cancelled() {
|
||||
output_queue.push(handler.event_name, error);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn synchronous_handlers_after_serialization_failure(
|
||||
output_queue: &AsyncHookOutputQueue,
|
||||
handlers: Vec<ConfiguredHandler>,
|
||||
error: &str,
|
||||
) -> Vec<ConfiguredHandler> {
|
||||
let (synchronous, asynchronous) = partition_handlers(handlers);
|
||||
spawn_async_failures(
|
||||
output_queue,
|
||||
asynchronous,
|
||||
format!("Async hook input serialization failed: {error}"),
|
||||
);
|
||||
synchronous
|
||||
}
|
||||
|
||||
pub(crate) fn running_summary(handler: &ConfiguredHandler) -> HookRunSummary {
|
||||
HookRunSummary {
|
||||
id: handler.run_id(),
|
||||
@@ -88,12 +176,16 @@ pub(crate) fn running_summary(handler: &ConfiguredHandler) -> HookRunSummary {
|
||||
|
||||
pub(crate) async fn execute_handlers<T>(
|
||||
shell: &CommandShell,
|
||||
output_queue: &AsyncHookOutputQueue,
|
||||
handlers: Vec<ConfiguredHandler>,
|
||||
input_json: String,
|
||||
cwd: &Path,
|
||||
turn_id: Option<String>,
|
||||
parse: fn(&ConfiguredHandler, CommandRunResult, Option<String>) -> ParsedHandler<T>,
|
||||
) -> Vec<ParsedHandler<T>> {
|
||||
let (handlers, asynchronous) = partition_handlers(handlers);
|
||||
spawn_async_handlers(shell, output_queue, asynchronous, input_json.clone(), cwd);
|
||||
|
||||
let mut pending = FuturesUnordered::new();
|
||||
for (configured_order, handler) in handlers.into_iter().enumerate() {
|
||||
let input_json = input_json.clone();
|
||||
@@ -175,6 +267,7 @@ mod tests {
|
||||
matcher: matcher.map(str::to_owned),
|
||||
command: command.to_string(),
|
||||
timeout_sec: 5,
|
||||
r#async: false,
|
||||
status_message: None,
|
||||
source_path: test_path_buf("/tmp/hooks.json").abs(),
|
||||
source: HookSource::User,
|
||||
@@ -416,7 +509,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn select_handlers_preserves_declaration_order() {
|
||||
let handlers = vec![
|
||||
let mut handlers = vec![
|
||||
make_handler(
|
||||
HookEventName::Stop,
|
||||
/*matcher*/ None,
|
||||
@@ -436,12 +529,35 @@ mod tests {
|
||||
/*display_order*/ 2,
|
||||
),
|
||||
];
|
||||
handlers[1].r#async = true;
|
||||
|
||||
let selected = select_handlers(&handlers, HookEventName::Stop, /*matcher_input*/ None);
|
||||
let synchronous = super::select_sync_handlers(
|
||||
&handlers,
|
||||
HookEventName::Stop,
|
||||
/*matcher_input*/ None,
|
||||
);
|
||||
let (_, asynchronous) = super::partition_handlers(selected.clone());
|
||||
|
||||
assert_eq!(selected.len(), 3);
|
||||
assert_eq!(selected[0].command, "first");
|
||||
assert_eq!(selected[1].command, "second");
|
||||
assert_eq!(selected[2].command, "third");
|
||||
assert_eq!(
|
||||
(
|
||||
synchronous
|
||||
.into_iter()
|
||||
.map(|handler| handler.command)
|
||||
.collect::<Vec<_>>(),
|
||||
asynchronous
|
||||
.into_iter()
|
||||
.map(|handler| handler.command)
|
||||
.collect::<Vec<_>>(),
|
||||
),
|
||||
(
|
||||
vec!["first".to_string(), "third".to_string()],
|
||||
vec!["second".to_string()],
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
pub(crate) mod async_output;
|
||||
pub(crate) mod command_runner;
|
||||
pub(crate) mod discovery;
|
||||
pub(crate) mod dispatcher;
|
||||
@@ -32,6 +33,9 @@ use codex_protocol::protocol::HookTrustStatus;
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use std::collections::HashMap;
|
||||
|
||||
pub use async_output::AsyncHookOutputBatch;
|
||||
pub use async_output::AsyncHookOutputQueue;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct CommandShell {
|
||||
pub program: String,
|
||||
@@ -44,6 +48,7 @@ pub(crate) struct ConfiguredHandler {
|
||||
pub matcher: Option<String>,
|
||||
pub command: String,
|
||||
pub timeout_sec: u64,
|
||||
pub r#async: bool,
|
||||
pub status_message: Option<String>,
|
||||
pub source_path: AbsolutePathBuf,
|
||||
pub source: HookSource,
|
||||
@@ -101,10 +106,12 @@ pub(crate) struct ClaudeHooksEngine {
|
||||
handlers: Vec<ConfiguredHandler>,
|
||||
warnings: Vec<String>,
|
||||
shell: CommandShell,
|
||||
async_output_queue: AsyncHookOutputQueue,
|
||||
output_spiller: HookOutputSpiller,
|
||||
}
|
||||
|
||||
impl ClaudeHooksEngine {
|
||||
#[cfg(test)]
|
||||
pub(crate) fn new(
|
||||
enabled: bool,
|
||||
bypass_hook_trust: bool,
|
||||
@@ -112,12 +119,33 @@ impl ClaudeHooksEngine {
|
||||
plugin_hook_sources: Vec<PluginHookSource>,
|
||||
plugin_hook_load_warnings: Vec<String>,
|
||||
shell: CommandShell,
|
||||
) -> Self {
|
||||
Self::new_with_async_output_queue(
|
||||
enabled,
|
||||
bypass_hook_trust,
|
||||
config_layer_stack,
|
||||
plugin_hook_sources,
|
||||
plugin_hook_load_warnings,
|
||||
shell,
|
||||
AsyncHookOutputQueue::default(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn new_with_async_output_queue(
|
||||
enabled: bool,
|
||||
bypass_hook_trust: bool,
|
||||
config_layer_stack: Option<&ConfigLayerStack>,
|
||||
plugin_hook_sources: Vec<PluginHookSource>,
|
||||
plugin_hook_load_warnings: Vec<String>,
|
||||
shell: CommandShell,
|
||||
async_output_queue: AsyncHookOutputQueue,
|
||||
) -> Self {
|
||||
if !enabled {
|
||||
return Self {
|
||||
handlers: Vec::new(),
|
||||
warnings: Vec::new(),
|
||||
shell,
|
||||
async_output_queue,
|
||||
output_spiller: HookOutputSpiller::new(),
|
||||
};
|
||||
}
|
||||
@@ -133,6 +161,7 @@ impl ClaudeHooksEngine {
|
||||
handlers: discovered.handlers,
|
||||
warnings: discovered.warnings,
|
||||
shell,
|
||||
async_output_queue,
|
||||
output_spiller: HookOutputSpiller::new(),
|
||||
}
|
||||
}
|
||||
@@ -172,8 +201,14 @@ impl ClaudeHooksEngine {
|
||||
turn_id: Option<String>,
|
||||
) -> SessionStartOutcome {
|
||||
let session_id = request.session_id;
|
||||
let mut outcome =
|
||||
crate::events::session_start::run(&self.handlers, &self.shell, request, turn_id).await;
|
||||
let mut outcome = crate::events::session_start::run(
|
||||
&self.handlers,
|
||||
&self.shell,
|
||||
&self.async_output_queue,
|
||||
request,
|
||||
turn_id,
|
||||
)
|
||||
.await;
|
||||
outcome.additional_contexts = self
|
||||
.maybe_spill_texts(session_id, outcome.additional_contexts)
|
||||
.await;
|
||||
@@ -182,8 +217,13 @@ impl ClaudeHooksEngine {
|
||||
|
||||
pub(crate) async fn run_pre_tool_use(&self, request: PreToolUseRequest) -> PreToolUseOutcome {
|
||||
let session_id = request.session_id;
|
||||
let mut outcome =
|
||||
crate::events::pre_tool_use::run(&self.handlers, &self.shell, request).await;
|
||||
let mut outcome = crate::events::pre_tool_use::run(
|
||||
&self.handlers,
|
||||
&self.shell,
|
||||
&self.async_output_queue,
|
||||
request,
|
||||
)
|
||||
.await;
|
||||
outcome.additional_contexts = self
|
||||
.maybe_spill_texts(session_id, outcome.additional_contexts)
|
||||
.await;
|
||||
@@ -194,7 +234,13 @@ impl ClaudeHooksEngine {
|
||||
&self,
|
||||
request: PermissionRequestRequest,
|
||||
) -> PermissionRequestOutcome {
|
||||
crate::events::permission_request::run(&self.handlers, &self.shell, request).await
|
||||
crate::events::permission_request::run(
|
||||
&self.handlers,
|
||||
&self.shell,
|
||||
&self.async_output_queue,
|
||||
request,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn run_post_tool_use(
|
||||
@@ -202,8 +248,13 @@ impl ClaudeHooksEngine {
|
||||
request: PostToolUseRequest,
|
||||
) -> PostToolUseOutcome {
|
||||
let session_id = request.session_id;
|
||||
let mut outcome =
|
||||
crate::events::post_tool_use::run(&self.handlers, &self.shell, request).await;
|
||||
let mut outcome = crate::events::post_tool_use::run(
|
||||
&self.handlers,
|
||||
&self.shell,
|
||||
&self.async_output_queue,
|
||||
request,
|
||||
)
|
||||
.await;
|
||||
outcome.additional_contexts = self
|
||||
.maybe_spill_texts(session_id, outcome.additional_contexts)
|
||||
.await;
|
||||
@@ -218,7 +269,13 @@ impl ClaudeHooksEngine {
|
||||
}
|
||||
|
||||
pub(crate) async fn run_pre_compact(&self, request: PreCompactRequest) -> PreCompactOutcome {
|
||||
crate::events::compact::run_pre(&self.handlers, &self.shell, request).await
|
||||
crate::events::compact::run_pre(
|
||||
&self.handlers,
|
||||
&self.shell,
|
||||
&self.async_output_queue,
|
||||
request,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn preview_post_compact(&self, request: &PostCompactRequest) -> Vec<HookRunSummary> {
|
||||
@@ -229,7 +286,13 @@ impl ClaudeHooksEngine {
|
||||
&self,
|
||||
request: PostCompactRequest,
|
||||
) -> StatelessHookOutcome {
|
||||
crate::events::compact::run_post(&self.handlers, &self.shell, request).await
|
||||
crate::events::compact::run_post(
|
||||
&self.handlers,
|
||||
&self.shell,
|
||||
&self.async_output_queue,
|
||||
request,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn preview_user_prompt_submit(
|
||||
@@ -244,8 +307,13 @@ impl ClaudeHooksEngine {
|
||||
request: UserPromptSubmitRequest,
|
||||
) -> UserPromptSubmitOutcome {
|
||||
let session_id = request.session_id;
|
||||
let mut outcome =
|
||||
crate::events::user_prompt_submit::run(&self.handlers, &self.shell, request).await;
|
||||
let mut outcome = crate::events::user_prompt_submit::run(
|
||||
&self.handlers,
|
||||
&self.shell,
|
||||
&self.async_output_queue,
|
||||
request,
|
||||
)
|
||||
.await;
|
||||
outcome.additional_contexts = self
|
||||
.maybe_spill_texts(session_id, outcome.additional_contexts)
|
||||
.await;
|
||||
@@ -258,7 +326,13 @@ impl ClaudeHooksEngine {
|
||||
|
||||
pub(crate) async fn run_stop(&self, request: StopRequest) -> StopOutcome {
|
||||
let session_id = request.session_id;
|
||||
let mut outcome = crate::events::stop::run(&self.handlers, &self.shell, request).await;
|
||||
let mut outcome = crate::events::stop::run(
|
||||
&self.handlers,
|
||||
&self.shell,
|
||||
&self.async_output_queue,
|
||||
request,
|
||||
)
|
||||
.await;
|
||||
outcome.continuation_fragments = self
|
||||
.maybe_spill_prompt_fragments(session_id, outcome.continuation_fragments)
|
||||
.await;
|
||||
|
||||
@@ -221,6 +221,7 @@ with Path(r"{log_path}").open("a", encoding="utf-8") as handle:
|
||||
plugin_hook_load_warnings: Vec::new(),
|
||||
shell_program: None,
|
||||
shell_args: Vec::new(),
|
||||
async_output_queue: Default::default(),
|
||||
});
|
||||
assert!(listed.hooks[0].is_managed);
|
||||
let cwd = cwd();
|
||||
@@ -1219,6 +1220,7 @@ print(json.dumps({
|
||||
plugin_hook_load_warnings: Vec::new(),
|
||||
shell_program: None,
|
||||
shell_args: Vec::new(),
|
||||
async_output_queue: Default::default(),
|
||||
});
|
||||
assert_eq!(
|
||||
listed.hooks[0].plugin_id.as_deref(),
|
||||
|
||||
Reference in New Issue
Block a user