openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Simplify async hook implementation and tests

Browse Source
This commit is contained in:
Abhinav Vedmala
2026-06-10 17:07:32 -07:00
parent d57e752091
commit 302b7f6a28
4 changed files with 188 additions and 351 deletions
Download Patch File Download Diff File Expand all files Collapse all files

313
codex-rs/core/tests/suite/hooks.rs
View File

@@ -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<TestCommandHook>) -> 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(()));

108
codex-rs/hooks/src/engine/async_command_tests.rs
View File

@@ -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());

80
codex-rs/hooks/src/engine/output_parser.rs
View File

@@ -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<String>,
) -> AsyncInformationalOutput {
AsyncInformationalOutput {
additional_context,
system_message: universal.system_message,
}
}
pub(crate) fn parse_subagent_stop(stdout: &str) -> Option<StopOutput> {
let wire: SubagentStopCommandOutputWire = parse_json(stdout)?;
Some(stop_output(

38
codex-rs/hooks/src/registry.rs
View File

@@ -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,