From 302b7f6a2823ca348f2dcf9a244b7fc696955fbc Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Wed, 10 Jun 2026 17:07:32 -0700 Subject: [PATCH] Simplify async hook implementation and tests --- codex-rs/core/tests/suite/hooks.rs | 313 ++++++------------ .../hooks/src/engine/async_command_tests.rs | 108 +++--- codex-rs/hooks/src/engine/output_parser.rs | 80 ++--- codex-rs/hooks/src/registry.rs | 38 +-- 4 files changed, 188 insertions(+), 351 deletions(-) diff --git a/codex-rs/core/tests/suite/hooks.rs b/codex-rs/core/tests/suite/hooks.rs index 94f5617fbd..18fccffffa 100644 --- a/codex-rs/core/tests/suite/hooks.rs +++ b/codex-rs/core/tests/suite/hooks.rs @@ -25,6 +25,7 @@ use codex_utils_absolute_path::AbsolutePathBuf; use core_test_support::hooks::trust_discovered_hooks; use core_test_support::hooks::trust_hooks; use core_test_support::managed_network_requirements_loader; +use core_test_support::responses::ResponseMock; use core_test_support::responses::ev_apply_patch_custom_tool_call; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -302,9 +303,38 @@ with Path(r"{log_path}").open("a", encoding="utf-8") as handle: Ok(()) } +struct TestCommandHook { + filename: &'static str, + source: String, + asynchronous: bool, +} + +fn write_command_hooks(home: &Path, event_name: &str, scripts: Vec) -> Result<()> { + let mut handlers = Vec::with_capacity(scripts.len()); + for script in scripts { + let script_path = home.join(script.filename); + fs::write(&script_path, script.source) + .with_context(|| format!("write test hook {}", script_path.display()))?; + handlers.push(serde_json::json!({ + "type": "command", + "command": format!("python3 {}", script_path.display()), + "async": script.asynchronous, + })); + } + let mut events = serde_json::Map::new(); + events.insert( + event_name.to_string(), + serde_json::json!([{ "hooks": handlers }]), + ); + fs::write( + home.join("hooks.json"), + serde_json::json!({ "hooks": events }).to_string(), + ) + .context("write hooks.json")?; + Ok(()) +} + fn write_async_user_prompt_submit_hooks(home: &Path) -> Result<()> { - let blocking_script_path = home.join("blocking_user_prompt_submit_hook.py"); - let async_script_path = home.join("async_user_prompt_submit_hook.py"); let ready_path = home.join("async_user_prompt_submit_ready"); let blocking_script = r#"import json import sys @@ -312,7 +342,8 @@ import sys payload = json.load(sys.stdin) if payload.get("prompt") == "blocked first prompt": print(json.dumps({"decision": "block", "reason": "blocked synchronously"})) -"#; +"# + .to_string(); let async_script = format!( r#"import json from pathlib import Path @@ -321,11 +352,6 @@ import sys payload = json.load(sys.stdin) prompt = payload.get("prompt") print(json.dumps({{ - "continue": False, - "stopReason": "ignored", - "systemMessage": f"async message for {{prompt}}", - "decision": "block", - "reason": "ignored", "hookSpecificOutput": {{ "hookEventName": "UserPromptSubmit", "additionalContext": f"async context for {{prompt}}" @@ -335,33 +361,25 @@ Path(r"{ready_path}").write_text(prompt, encoding="utf-8") "#, ready_path = ready_path.display(), ); - let hooks = serde_json::json!({ - "hooks": { - "UserPromptSubmit": [{ - "hooks": [ - { - "type": "command", - "command": format!("python3 {}", blocking_script_path.display()), - }, - { - "type": "command", - "command": format!("python3 {}", async_script_path.display()), - "async": true, - } - ] - }] - } - }); - - fs::write(blocking_script_path, blocking_script) - .context("write blocking user prompt submit hook")?; - fs::write(async_script_path, async_script).context("write async user prompt submit hook")?; - fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; - Ok(()) + write_command_hooks( + home, + "UserPromptSubmit", + vec![ + TestCommandHook { + filename: "blocking_user_prompt_submit_hook.py", + source: blocking_script, + asynchronous: false, + }, + TestCommandHook { + filename: "async_user_prompt_submit_hook.py", + source: async_script, + asynchronous: true, + }, + ], + ) } fn write_gated_async_user_prompt_submit_hook(home: &Path) -> Result<()> { - let script_path = home.join("gated_async_user_prompt_submit_hook.py"); let started_path = home.join("gated_async_user_prompt_submit_started"); let release_path = home.join("gated_async_user_prompt_submit_release"); let ready_path = home.join("gated_async_user_prompt_submit_ready"); @@ -389,61 +407,18 @@ Path(r"{ready_path}").write_text(prompt, encoding="utf-8") release_path = release_path.display(), ready_path = ready_path.display(), ); - let hooks = serde_json::json!({ - "hooks": { - "UserPromptSubmit": [{ - "hooks": [{ - "type": "command", - "command": format!("python3 {}", script_path.display()), - "async": true, - }] - }] - } - }); - - fs::write(script_path, script).context("write gated async user prompt submit hook")?; - fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; - Ok(()) -} - -fn write_async_session_start_hook_with_system_message( - home: &Path, - system_message: &str, -) -> Result<()> { - let script_path = home.join("async_session_start_system_message_hook.py"); - let ready_path = home.join("async_session_start_system_message_ready"); - let system_message_json = - serde_json::to_string(system_message).context("serialize async system message")?; - let script = format!( - r#"import json -from pathlib import Path -import sys - -json.load(sys.stdin) -print(json.dumps({{"systemMessage": {system_message_json}}})) -Path(r"{ready_path}").write_text("ready", encoding="utf-8") -"#, - ready_path = ready_path.display(), - ); - let hooks = serde_json::json!({ - "hooks": { - "SessionStart": [{ - "hooks": [{ - "type": "command", - "command": format!("python3 {}", script_path.display()), - "async": true, - }] - }] - } - }); - - fs::write(script_path, script).context("write async session start system message hook")?; - fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; - Ok(()) + write_command_hooks( + home, + "UserPromptSubmit", + vec![TestCommandHook { + filename: "gated_async_user_prompt_submit_hook.py", + source: script, + asynchronous: true, + }], + ) } fn write_async_session_start_hook(home: &Path) -> Result<()> { - let script_path = home.join("async_session_start_hook.py"); let ready_path = home.join("async_session_start_ready"); let script = format!( r#"import json @@ -452,8 +427,6 @@ import sys json.load(sys.stdin) print(json.dumps({{ - "continue": False, - "stopReason": "ignored", "systemMessage": "async startup message", "hookSpecificOutput": {{ "hookEventName": "SessionStart", @@ -464,21 +437,15 @@ Path(r"{ready_path}").write_text("ready", encoding="utf-8") "#, ready_path = ready_path.display(), ); - let hooks = serde_json::json!({ - "hooks": { - "SessionStart": [{ - "hooks": [{ - "type": "command", - "command": format!("python3 {}", script_path.display()), - "async": true, - }] - }] - } - }); - - fs::write(script_path, script).context("write async session start hook")?; - fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; - Ok(()) + write_command_hooks( + home, + "SessionStart", + vec![TestCommandHook { + filename: "async_session_start_hook.py", + source: script, + asynchronous: true, + }], + ) } async fn wait_for_async_hook(path: &Path) -> Result<()> { @@ -493,6 +460,25 @@ async fn wait_for_async_hook(path: &Path) -> Result<()> { Ok(()) } +async fn mount_two_turn_responses(server: &wiremock::MockServer) -> ResponseMock { + mount_sse_sequence( + server, + ["first", "second"] + .into_iter() + .enumerate() + .map(|(index, message)| { + let number = index + 1; + sse(vec![ + ev_response_created(&format!("resp-{number}")), + ev_assistant_message(&format!("msg-{number}"), message), + ev_completed(&format!("resp-{number}")), + ]) + }) + .collect(), + ) + .await +} + fn write_pre_tool_use_hook( home: &Path, matcher: Option<&str>, @@ -1600,22 +1586,7 @@ async fn async_startup_output_skips_first_accepted_turn() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let responses = mount_sse_sequence( - &server, - vec![ - sse(vec![ - ev_response_created("resp-1"), - ev_assistant_message("msg-1", "first"), - ev_completed("resp-1"), - ]), - sse(vec![ - ev_response_created("resp-2"), - ev_assistant_message("msg-2", "second"), - ev_completed("resp-2"), - ]), - ], - ) - .await; + let responses = mount_two_turn_responses(&server).await; let mut builder = test_codex() .with_pre_build_hook(|home| { @@ -1626,7 +1597,27 @@ async fn async_startup_output_skips_first_accepted_turn() -> Result<()> { test.submit_turn("first prompt").await?; wait_for_async_hook(&test.codex_home_path().join("async_session_start_ready")).await?; - test.submit_turn("second prompt").await?; + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "second prompt".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + let warning = wait_for_event(&test.codex, |event| matches!(event, EventMsg::Warning(_))).await; + let EventMsg::Warning(warning) = warning else { + unreachable!("waited for warning event") + }; + assert_eq!(warning.message, "async startup message"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; let requests = responses.requests(); assert_eq!(requests.len(), 2); @@ -1649,22 +1640,7 @@ async fn async_output_survives_runtime_config_refresh() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let responses = mount_sse_sequence( - &server, - vec![ - sse(vec![ - ev_response_created("resp-1"), - ev_assistant_message("msg-1", "first"), - ev_completed("resp-1"), - ]), - sse(vec![ - ev_response_created("resp-2"), - ev_assistant_message("msg-2", "second"), - ev_completed("resp-2"), - ]), - ], - ) - .await; + let responses = mount_two_turn_responses(&server).await; let mut builder = test_codex() .with_pre_build_hook(|home| { @@ -1708,75 +1684,6 @@ async fn async_output_survives_runtime_config_refresh() -> Result<()> { Ok(()) } -#[tokio::test] -async fn async_system_message_spills_large_output() -> Result<()> { - skip_if_no_network!(Ok(())); - - let server = start_mock_server().await; - mount_sse_sequence( - &server, - vec![ - sse(vec![ - ev_response_created("resp-1"), - ev_assistant_message("msg-1", "first"), - ev_completed("resp-1"), - ]), - sse(vec![ - ev_response_created("resp-2"), - ev_assistant_message("msg-2", "second"), - ev_completed("resp-2"), - ]), - ], - ) - .await; - let system_message = "async warning output ".repeat(800); - - let mut builder = test_codex() - .with_pre_build_hook({ - let system_message = system_message.clone(); - move |home| { - write_async_session_start_hook_with_system_message(home, &system_message) - .expect("write async session start system message hook"); - } - }) - .with_config(trust_discovered_hooks); - let test = builder.build(&server).await?; - - test.submit_turn("first prompt").await?; - wait_for_async_hook( - &test - .codex_home_path() - .join("async_session_start_system_message_ready"), - ) - .await?; - test.codex - .submit(Op::UserInput { - items: vec![UserInput::Text { - text: "second prompt".to_string(), - text_elements: Vec::new(), - }], - final_output_json_schema: None, - responsesapi_client_metadata: None, - additional_context: Default::default(), - thread_settings: Default::default(), - }) - .await?; - - let warning = wait_for_event(&test.codex, |event| matches!(event, EventMsg::Warning(_))).await; - let EventMsg::Warning(warning) = warning else { - unreachable!("waited for warning event") - }; - assert!(warning.message.contains("tokens truncated")); - let path = spilled_hook_output_path(&warning.message).context("spilled system message path")?; - assert_eq!(fs::read_to_string(path)?, system_message); - wait_for_event(&test.codex, |event| { - matches!(event, EventMsg::TurnComplete(_)) - }) - .await; - - Ok(()) -} - #[tokio::test] async fn session_start_hook_spills_large_additional_context() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/hooks/src/engine/async_command_tests.rs b/codex-rs/hooks/src/engine/async_command_tests.rs index d4c516ad1e..a0e13aebce 100644 --- a/codex-rs/hooks/src/engine/async_command_tests.rs +++ b/codex-rs/hooks/src/engine/async_command_tests.rs @@ -13,42 +13,31 @@ use codex_protocol::protocol::HookSource; use codex_utils_absolute_path::AbsolutePathBuf; use std::collections::HashMap; +enum TestOutput<'a> { + AdditionalContext(&'a str), + SystemMessage(&'a str), +} + fn complete( runtime: &AsyncCommandRuntime, launch_sequence: u64, deliver_at_generation: u64, - text: &str, + output: TestOutput<'_>, ) { let mut state = runtime.inner.state.lock().expect("async hook state"); let ready_sequence = state.next_ready_sequence; state.next_ready_sequence += 1; + let (additional_context, system_message) = match output { + TestOutput::AdditionalContext(text) => (Some(text.to_string()), None), + TestOutput::SystemMessage(text) => (None, Some(text.to_string())), + }; state.completions.insert( launch_sequence, AsyncHookCompletion { deliver_at_generation, ready_sequence, - additional_context: Some(text.to_string()), - system_message: None, - }, - ); -} - -fn complete_with_system_message( - runtime: &AsyncCommandRuntime, - launch_sequence: u64, - deliver_at_generation: u64, - text: &str, -) { - let mut state = runtime.inner.state.lock().expect("async hook state"); - let ready_sequence = state.next_ready_sequence; - state.next_ready_sequence += 1; - state.completions.insert( - launch_sequence, - AsyncHookCompletion { - deliver_at_generation, - ready_sequence, - additional_context: None, - system_message: Some(text.to_string()), + additional_context, + system_message, }, ); } @@ -58,7 +47,10 @@ fn completion_after_cutoff_waits_for_following_accepted_turn() { let runtime = AsyncCommandRuntime::new(); let cutoff = runtime.delivery_cutoff(); complete( - &runtime, /*launch_sequence*/ 0, /*deliver_at_generation*/ 1, "late", + &runtime, + /*launch_sequence*/ 0, + /*deliver_at_generation*/ 1, + TestOutput::AdditionalContext("late"), ); assert_eq!( @@ -70,41 +62,14 @@ fn completion_after_cutoff_waits_for_following_accepted_turn() { assert_eq!(delivery.additional_contexts, vec!["late"]); } -#[test] -fn startup_completion_skips_first_accepted_turn() { - let runtime = AsyncCommandRuntime::new(); - complete( - &runtime, /*launch_sequence*/ 0, /*deliver_at_generation*/ 2, "startup", - ); - - assert_eq!( - runtime.commit_accepted_turn_and_drain(runtime.delivery_cutoff()), - Default::default() - ); - - let delivery = runtime.commit_accepted_turn_and_drain(runtime.delivery_cutoff()); - assert_eq!(delivery.additional_contexts, vec!["startup"]); -} - -#[test] -fn blocked_submission_does_not_advance_generation() { - let runtime = AsyncCommandRuntime::new(); - complete( - &runtime, - /*launch_sequence*/ 0, - /*deliver_at_generation*/ 1, - "after block", - ); - - let delivery = runtime.commit_accepted_turn_and_drain(runtime.delivery_cutoff()); - assert_eq!(delivery.additional_contexts, vec!["after block"]); -} - #[test] fn unfinished_earlier_launch_does_not_block_ready_output() { let runtime = AsyncCommandRuntime::new(); complete( - &runtime, /*launch_sequence*/ 1, /*deliver_at_generation*/ 1, "ready", + &runtime, + /*launch_sequence*/ 1, + /*deliver_at_generation*/ 1, + TestOutput::AdditionalContext("ready"), ); let delivery = runtime.commit_accepted_turn_and_drain(runtime.delivery_cutoff()); @@ -158,19 +123,34 @@ fn shared_output_budget_leaves_remaining_completions_queued() { let runtime = AsyncCommandRuntime::new(); let output = "x".repeat(MAX_DELIVERED_OUTPUT_TOKENS_PER_TURN); complete( - &runtime, /*launch_sequence*/ 0, /*deliver_at_generation*/ 1, &output, - ); - complete_with_system_message( - &runtime, /*launch_sequence*/ 1, /*deliver_at_generation*/ 1, &output, + &runtime, + /*launch_sequence*/ 0, + /*deliver_at_generation*/ 1, + TestOutput::AdditionalContext(&output), ); complete( - &runtime, /*launch_sequence*/ 2, /*deliver_at_generation*/ 1, &output, - ); - complete_with_system_message( - &runtime, /*launch_sequence*/ 3, /*deliver_at_generation*/ 1, &output, + &runtime, + /*launch_sequence*/ 1, + /*deliver_at_generation*/ 1, + TestOutput::SystemMessage(&output), ); complete( - &runtime, /*launch_sequence*/ 4, /*deliver_at_generation*/ 1, &output, + &runtime, + /*launch_sequence*/ 2, + /*deliver_at_generation*/ 1, + TestOutput::AdditionalContext(&output), + ); + complete( + &runtime, + /*launch_sequence*/ 3, + /*deliver_at_generation*/ 1, + TestOutput::SystemMessage(&output), + ); + complete( + &runtime, + /*launch_sequence*/ 4, + /*deliver_at_generation*/ 1, + TestOutput::AdditionalContext(&output), ); let first_delivery = runtime.commit_accepted_turn_and_drain(runtime.delivery_cutoff()); diff --git a/codex-rs/hooks/src/engine/output_parser.rs b/codex-rs/hooks/src/engine/output_parser.rs index 4ece51f1db..01ce9c9b83 100644 --- a/codex-rs/hooks/src/engine/output_parser.rs +++ b/codex-rs/hooks/src/engine/output_parser.rs @@ -320,64 +320,30 @@ pub(crate) fn parse_async_informational( } let parsed = match event_name { - HookEventName::SessionStart => { - parse_session_start(stdout).map(|output| AsyncInformationalOutput { - additional_context: output.additional_context, - system_message: output.universal.system_message, - }) - } - HookEventName::SubagentStart => { - parse_subagent_start(stdout).map(|output| AsyncInformationalOutput { - additional_context: output.additional_context, - system_message: output.universal.system_message, - }) - } - HookEventName::PreToolUse => { - parse_pre_tool_use(stdout).map(|output| AsyncInformationalOutput { - additional_context: output.additional_context, - system_message: output.universal.system_message, - }) - } - HookEventName::PermissionRequest => { - parse_permission_request(stdout).map(|output| AsyncInformationalOutput { - additional_context: None, - system_message: output.universal.system_message, - }) - } - HookEventName::PostToolUse => { - parse_post_tool_use(stdout).map(|output| AsyncInformationalOutput { - additional_context: output.additional_context, - system_message: output.universal.system_message, - }) - } + HookEventName::SessionStart => parse_session_start(stdout) + .map(|output| async_informational(output.universal, output.additional_context)), + HookEventName::SubagentStart => parse_subagent_start(stdout) + .map(|output| async_informational(output.universal, output.additional_context)), + HookEventName::PreToolUse => parse_pre_tool_use(stdout) + .map(|output| async_informational(output.universal, output.additional_context)), + HookEventName::PermissionRequest => parse_permission_request(stdout) + .map(|output| async_informational(output.universal, None)), + HookEventName::PostToolUse => parse_post_tool_use(stdout) + .map(|output| async_informational(output.universal, output.additional_context)), HookEventName::PreCompact => { - parse_pre_compact(stdout).map(|output| AsyncInformationalOutput { - additional_context: None, - system_message: output.universal.system_message, - }) + parse_pre_compact(stdout).map(|output| async_informational(output.universal, None)) } HookEventName::PostCompact => { - parse_post_compact(stdout).map(|output| AsyncInformationalOutput { - additional_context: None, - system_message: output.universal.system_message, - }) - } - HookEventName::UserPromptSubmit => { - parse_user_prompt_submit(stdout).map(|output| AsyncInformationalOutput { - additional_context: output.additional_context, - system_message: output.universal.system_message, - }) + parse_post_compact(stdout).map(|output| async_informational(output.universal, None)) } + HookEventName::UserPromptSubmit => parse_user_prompt_submit(stdout) + .map(|output| async_informational(output.universal, output.additional_context)), HookEventName::SubagentStop => { - parse_subagent_stop(stdout).map(|output| AsyncInformationalOutput { - additional_context: None, - system_message: output.universal.system_message, - }) + parse_subagent_stop(stdout).map(|output| async_informational(output.universal, None)) + } + HookEventName::Stop => { + parse_stop(stdout).map(|output| async_informational(output.universal, None)) } - HookEventName::Stop => parse_stop(stdout).map(|output| AsyncInformationalOutput { - additional_context: None, - system_message: output.universal.system_message, - }), }; if let Some(parsed) = parsed { return (parsed.additional_context.is_some() || parsed.system_message.is_some()) @@ -399,6 +365,16 @@ pub(crate) fn parse_async_informational( }) } +fn async_informational( + universal: UniversalOutput, + additional_context: Option, +) -> AsyncInformationalOutput { + AsyncInformationalOutput { + additional_context, + system_message: universal.system_message, + } +} + pub(crate) fn parse_subagent_stop(stdout: &str) -> Option { let wire: SubagentStopCommandOutputWire = parse_json(stdout)?; Some(stop_output( diff --git a/codex-rs/hooks/src/registry.rs b/codex-rs/hooks/src/registry.rs index 6c99889a56..f727c7aebb 100644 --- a/codex-rs/hooks/src/registry.rs +++ b/codex-rs/hooks/src/registry.rs @@ -61,40 +61,14 @@ impl Default for Hooks { impl Hooks { pub fn new(config: HooksConfig) -> Self { - let HooksConfig { - legacy_notify_argv, - feature_enabled, - bypass_hook_trust, - config_layer_stack, - plugin_hook_sources, - plugin_hook_load_warnings, - shell_program, - shell_args, - } = config; - let after_agent = legacy_notify_argv - .filter(|argv| !argv.is_empty() && !argv[0].is_empty()) - .map(crate::notify_hook) - .into_iter() - .collect(); - let engine = ClaudeHooksEngine::new( - feature_enabled, - bypass_hook_trust, - config_layer_stack.as_ref(), - plugin_hook_sources, - plugin_hook_load_warnings, - CommandShell { - program: shell_program.unwrap_or_default(), - args: shell_args, - }, - AsyncCommandRuntime::new(), - ); - Self { - after_agent, - engine, - } + Self::from_config(config, AsyncCommandRuntime::new()) } pub fn reconfigured(&self, config: HooksConfig) -> Self { + Self::from_config(config, self.engine.async_runtime()) + } + + fn from_config(config: HooksConfig, async_runtime: AsyncCommandRuntime) -> Self { let HooksConfig { legacy_notify_argv, feature_enabled, @@ -120,7 +94,7 @@ impl Hooks { program: shell_program.unwrap_or_default(), args: shell_args, }, - self.engine.async_runtime(), + async_runtime, ); Self { after_agent,