Finish async hook follow-ups

This commit is contained in:
Abhinav Vedmala
2026-06-10 16:37:21 -07:00
parent 7960750d8b
commit d57e752091
13 changed files with 365 additions and 34 deletions
+215
View File
@@ -360,6 +360,88 @@ Path(r"{ready_path}").write_text(prompt, encoding="utf-8")
Ok(())
}
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");
let script = format!(
r#"import json
from pathlib import Path
import sys
import time
payload = json.load(sys.stdin)
prompt = payload.get("prompt")
Path(r"{started_path}").write_text(prompt, encoding="utf-8")
release_path = Path(r"{release_path}")
while not release_path.exists():
time.sleep(0.01)
print(json.dumps({{
"hookSpecificOutput": {{
"hookEventName": "UserPromptSubmit",
"additionalContext": f"async context surviving refresh for {{prompt}}"
}}
}}))
Path(r"{ready_path}").write_text(prompt, encoding="utf-8")
"#,
started_path = started_path.display(),
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(())
}
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");
@@ -1562,6 +1644,139 @@ async fn async_startup_output_skips_first_accepted_turn() -> Result<()> {
Ok(())
}
#[tokio::test]
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 mut builder = test_codex()
.with_pre_build_hook(|home| {
write_gated_async_user_prompt_submit_hook(home)
.expect("write gated async user prompt submit hook");
})
.with_config(trust_discovered_hooks);
let test = builder.build(&server).await?;
test.submit_turn("before refresh").await?;
wait_for_async_hook(
&test
.codex_home_path()
.join("gated_async_user_prompt_submit_started"),
)
.await?;
test.codex.refresh_runtime_config(test.config.clone()).await;
fs::write(
test.codex_home_path()
.join("gated_async_user_prompt_submit_release"),
"release",
)?;
wait_for_async_hook(
&test
.codex_home_path()
.join("gated_async_user_prompt_submit_ready"),
)
.await?;
test.submit_turn("after refresh").await?;
let requests = responses.requests();
assert_eq!(requests.len(), 2);
assert!(
requests[1]
.message_input_texts("developer")
.contains(&"async context surviving refresh for before refresh".to_string())
);
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(()));
+22 -7
View File
@@ -37,7 +37,7 @@ use crate::output_spill::HookOutputSpiller;
const MAX_QUEUED_COMPLETIONS: usize = 64;
const MAX_IN_FLIGHT_COMMANDS: usize = 32;
const MAX_DELIVERED_COMPLETIONS_PER_TURN: usize = 8;
const MAX_DELIVERED_CONTEXT_TOKENS_PER_TURN: usize = 10_000;
const MAX_DELIVERED_OUTPUT_TOKENS_PER_TURN: usize = 10_000;
/// Informational async hook output ready to be recorded for a model request.
#[derive(Debug, Default, PartialEq, Eq)]
@@ -209,6 +209,14 @@ impl AsyncCommandRuntime {
.await,
);
}
if let Some(system_message) = output.system_message.take() {
output.system_message = Some(
inner
.output_spiller
.maybe_spill_text(thread_id, system_message)
.await,
);
}
if inner.shutting_down.load(Ordering::Acquire) {
return;
}
@@ -265,7 +273,7 @@ impl AsyncCommandRuntime {
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut eligible = Vec::new();
let mut context_tokens = 0usize;
let mut output_tokens = 0usize;
for (launch_sequence, completion) in &state.completions {
if completion.ready_sequence >= cutoff.ready_sequence
|| completion.deliver_at_generation > accepted_generation
@@ -275,18 +283,25 @@ impl AsyncCommandRuntime {
if eligible.len() >= MAX_DELIVERED_COMPLETIONS_PER_TURN {
break;
}
let completion_context_tokens = completion
let completion_output_tokens = completion
.additional_context
.as_deref()
.map(approx_token_count)
.unwrap_or_default();
if context_tokens.saturating_add(completion_context_tokens)
> MAX_DELIVERED_CONTEXT_TOKENS_PER_TURN
.unwrap_or_default()
.saturating_add(
completion
.system_message
.as_deref()
.map(approx_token_count)
.unwrap_or_default(),
);
if output_tokens.saturating_add(completion_output_tokens)
> MAX_DELIVERED_OUTPUT_TOKENS_PER_TURN
{
break;
}
eligible.push(*launch_sequence);
context_tokens = context_tokens.saturating_add(completion_context_tokens);
output_tokens = output_tokens.saturating_add(completion_output_tokens);
}
let mut delivery = AsyncHookDelivery::default();
for launch_sequence in eligible {
@@ -2,7 +2,7 @@ use pretty_assertions::assert_eq;
use super::AsyncCommandRuntime;
use super::AsyncHookCompletion;
use super::MAX_DELIVERED_CONTEXT_TOKENS_PER_TURN;
use super::MAX_DELIVERED_OUTPUT_TOKENS_PER_TURN;
use super::MAX_IN_FLIGHT_COMMANDS;
use crate::engine::CommandShell;
use crate::engine::ConfiguredHandler;
@@ -33,6 +33,26 @@ fn complete(
);
}
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()),
},
);
}
#[test]
fn completion_after_cutoff_waits_for_following_accepted_turn() {
let runtime = AsyncCommandRuntime::new();
@@ -134,21 +154,34 @@ async fn launch_is_skipped_at_session_concurrency_limit() {
}
#[test]
fn context_delivery_budget_leaves_remaining_completions_queued() {
fn shared_output_budget_leaves_remaining_completions_queued() {
let runtime = AsyncCommandRuntime::new();
let context = "x".repeat(MAX_DELIVERED_CONTEXT_TOKENS_PER_TURN);
for launch_sequence in 0..5 {
complete(
&runtime,
launch_sequence,
/*deliver_at_generation*/ 1,
&context,
);
}
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,
);
complete(
&runtime, /*launch_sequence*/ 2, /*deliver_at_generation*/ 1, &output,
);
complete_with_system_message(
&runtime, /*launch_sequence*/ 3, /*deliver_at_generation*/ 1, &output,
);
complete(
&runtime, /*launch_sequence*/ 4, /*deliver_at_generation*/ 1, &output,
);
let first_delivery = runtime.commit_accepted_turn_and_drain(runtime.delivery_cutoff());
assert_eq!(first_delivery.additional_contexts, vec![context.clone(); 4]);
assert_eq!(
first_delivery,
super::AsyncHookDelivery {
additional_contexts: vec![output.clone(), output.clone()],
system_messages: vec![output.clone(), output.clone()],
}
);
let second_delivery = runtime.commit_accepted_turn_and_drain(runtime.delivery_cutoff());
assert_eq!(second_delivery.additional_contexts, vec![context]);
assert_eq!(second_delivery.additional_contexts, vec![output]);
}
+46 -14
View File
@@ -5,6 +5,7 @@ use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::approx_token_count;
use codex_utils_output_truncation::formatted_truncate_text;
use tokio::fs;
use tokio::io::AsyncWriteExt;
use tracing::warn;
use uuid::Uuid;
@@ -20,14 +21,14 @@ impl HookOutputSpiller {
pub(crate) fn new() -> Self {
Self {
output_dir: AbsolutePathBuf::resolve_path_against_base(std::env::temp_dir(), "/")
.join(HOOK_OUTPUTS_DIR),
.join(format!("{HOOK_OUTPUTS_DIR}-{}", Uuid::new_v4())),
}
}
/// Keeps hook text within the model-visible hook-output budget.
///
/// Oversized text is written in full under the OS temp directory at
/// `<temp_dir>/hook_outputs/<thread_id>/`
/// `<temp_dir>/hook_outputs-<runtime_id>/<thread_id>/`
/// and replaced with the same head/tail preview style used for other truncated
/// output, plus a path back to the preserved full text.
pub(crate) async fn maybe_spill_text(&self, thread_id: ThreadId, text: String) -> String {
@@ -36,20 +37,51 @@ impl HookOutputSpiller {
}
let path = hook_output_path(&self.output_dir, thread_id);
if let Some(parent) = path.parent()
&& let Err(err) = fs::create_dir_all(parent.as_ref()).await
{
warn!(
"failed to create hook output directory {}: {err}",
parent.display()
);
return formatted_truncate_text(
&text,
TruncationPolicy::Tokens(HOOK_OUTPUT_TOKEN_LIMIT),
);
if let Some(parent) = path.parent() {
let mut dir_builder = fs::DirBuilder::new();
dir_builder.recursive(true);
#[cfg(unix)]
dir_builder.mode(0o700);
if let Err(err) = dir_builder.create(parent.as_ref()).await {
warn!(
"failed to create hook output directory {}: {err}",
parent.display()
);
return formatted_truncate_text(
&text,
TruncationPolicy::Tokens(HOOK_OUTPUT_TOKEN_LIMIT),
);
}
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
if let Err(err) =
fs::set_permissions(parent.as_ref(), std::fs::Permissions::from_mode(0o700))
.await
{
warn!(
"failed to secure hook output directory {}: {err}",
parent.display()
);
return formatted_truncate_text(
&text,
TruncationPolicy::Tokens(HOOK_OUTPUT_TOKEN_LIMIT),
);
}
}
}
if let Err(err) = fs::write(path.as_ref(), &text).await {
let mut options = fs::OpenOptions::new();
options.write(true).create_new(true);
#[cfg(unix)]
options.mode(0o600);
let write_result = async {
let mut file = options.open(path.as_ref()).await?;
file.write_all(text.as_bytes()).await
}
.await;
if let Err(err) = write_result {
warn!("failed to write hook output {}: {err}", path.display());
return formatted_truncate_text(
&text,
+16
View File
@@ -3,6 +3,11 @@ use anyhow::Context;
use anyhow::Result;
use tempfile::tempdir;
#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;
#[cfg(unix)]
use std::path::Path;
#[tokio::test]
async fn small_hook_output_remains_inline() -> Result<()> {
let dir = tempdir()?;
@@ -38,5 +43,16 @@ async fn large_hook_output_spills_to_file() -> Result<()> {
.find_map(|line| line.strip_prefix("Full hook output saved to: "))
.context("spill path")?;
assert_eq!(fs::read_to_string(path).await?, text);
#[cfg(unix)]
{
let path = Path::new(path);
let parent_mode = fs::metadata(path.parent().context("spill parent")?)
.await?
.permissions()
.mode()
& 0o777;
let file_mode = fs::metadata(path).await?.permissions().mode() & 0o777;
assert_eq!((parent_mode, file_mode), (0o700, 0o600));
}
Ok(())
}
@@ -1,4 +1,5 @@
use codex_app_server_protocol::HookEventName;
use codex_app_server_protocol::HookExecutionMode;
use codex_app_server_protocol::HookMetadata;
use codex_app_server_protocol::HookSource;
use codex_app_server_protocol::HookTrustStatus;
@@ -491,6 +492,10 @@ impl HooksBrowserView {
width,
Some(MAX_COMMAND_DETAIL_LINES),
));
lines.push(detail_line(
"Mode",
hook_execution_mode_label(hook.execution_mode),
));
lines.push(detail_line("Timeout", &format!("{}s", hook.timeout_sec)));
lines.push(detail_line("Trust", hook_trust_label(hook.trust_status)));
lines
@@ -727,6 +732,13 @@ fn hook_trust_label(status: HookTrustStatus) -> &'static str {
}
}
fn hook_execution_mode_label(mode: HookExecutionMode) -> &'static str {
match mode {
HookExecutionMode::Sync => "Sync",
HookExecutionMode::Async => "Async",
}
}
fn event_label(event_name: HookEventName) -> &'static str {
match event_name {
HookEventName::PreToolUse => "PreToolUse",
@@ -1094,6 +1106,7 @@ mod tests {
/*display_order*/ 0,
);
untrusted_hook.trust_status = HookTrustStatus::Untrusted;
untrusted_hook.execution_mode = HookExecutionMode::Async;
let mut view = HooksBrowserView::new(
vec![untrusted_hook],
Vec::new(),
@@ -14,6 +14,7 @@ expression: "render_lines(&view, 44)"
Command one two three four five six
seven eight nine ten eleven
twelve thirteen fourteen…
Mode Sync
Timeout 30s
Trust Trusted
@@ -13,6 +13,7 @@ expression: "render_lines(&view, 112)"
Matcher Bash
Source Plugin - superpowers@openai-curated
Command ${CODEX_PLUGIN_ROOT}/hooks/pre-tool-use-check.sh
Mode Sync
Timeout 30s
Trust Trusted
@@ -12,6 +12,7 @@ expression: "render_lines(&view, 112)"
Matcher Bash
Source Admin config
Command /enterprise/hooks/permission-check.sh
Mode Sync
Timeout 30s
Trust Managed
@@ -12,6 +12,7 @@ expression: "render_lines(&view, 112)"
Matcher Bash
Source User config - /tmp/hooks.json
Command /tmp/pre-tool-use-check.sh
Mode Sync
Timeout 30s
Trust New hook - review required
@@ -19,6 +19,7 @@ expression: "render_lines(&view, 112)"
Matcher Bash
Source User config - /tmp/hooks.json
Command /tmp/hook-8.sh
Mode Sync
Timeout 30s
Trust Trusted
@@ -13,6 +13,7 @@ expression: "render_lines(&view, 112)"
Matcher Bash
Source Admin config
Command /enterprise/hooks/pre-tool-use-2.sh
Mode Sync
Timeout 30s
Trust Managed
@@ -12,6 +12,7 @@ expression: "render_lines(&view, 112)"
Matcher Bash
Source User config - /tmp/hooks.json
Command ~/bin/untrusted.sh
Mode Async
Timeout 30s
Trust New hook - review required