From 629d90398f020e2610542160f7c7e44c9a5bc55f Mon Sep 17 00:00:00 2001 From: Bernardo Kuri Date: Tue, 18 Aug 2026 19:34:30 -0600 Subject: [PATCH 1/8] feat: configurable timeout reply and timeout hook Add two optional config fields: timeout_reply overrides the hardcoded timeout message, timeout_hook runs a shell command on run timeout whose stdout becomes the reply (env: PUSH_THREAD, PUSH_ROW_ID, PUSH_BACKEND, PUSH_WORK_DIR; 5s budget; warn and fall back on any hook failure). FakeRunner now enforces the caller's run timeout like the real runners, enabling gateway-level timeout tests. --- docs/configuration.md | 2 + src/agent.rs | 11 +++- src/config.rs | 8 +++ src/gateway/tests.rs | 131 ++++++++++++++++++++++++++++++++++++++++++ src/gateway/worker.rs | 46 ++++++++++++++- src/test_support.rs | 2 + 6 files changed, 196 insertions(+), 4 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index d8207d4..fcdc233 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -209,6 +209,8 @@ requests. Review [permissions and security](security.md) before enabling jobs. | `agent` | `"claude"` | Default backend | | `poll_interval` | `"3s"` | Delay between channel polls | | `run_timeout` | `"10m"` | Maximum chat backend run time | +| `timeout_reply` | hardcoded string | Overrides the default timeout reply text | +| `timeout_hook` | none | Shell command run on run timeout; stdout becomes the reply (env vars: `PUSH_THREAD`, `PUSH_ROW_ID`, `PUSH_BACKEND`, `PUSH_WORK_DIR`; 5s budget, warn and fall back on any failure) | ### iMessage diff --git a/src/agent.rs b/src/agent.rs index 529fff5..67bed8d 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -192,7 +192,7 @@ pub struct FakeRunCall { #[cfg(test)] impl FakeRunner { - async fn run(&self, req: Request<'_>, _timeout: Duration) -> Result { + async fn run(&self, req: Request<'_>, timeout: Duration) -> Result { self.calls.lock().unwrap().push(FakeRunCall { session_id: req.session_id.to_string(), is_new: req.is_new, @@ -215,7 +215,14 @@ impl FakeRunner { before_return(); } if let Some(release) = &self.wait_for_release { - release.notified().await; + // Real runners enforce the caller's timeout themselves; so does the + // fake, so a release that never arrives surfaces as Timeout. + if tokio::time::timeout(timeout, release.notified()) + .await + .is_err() + { + return Err(RunError::Timeout); + } } if let Some(message) = &self.failure { return Err(RunError::Failed(message.clone())); diff --git a/src/config.rs b/src/config.rs index 0d075da..ced588a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -52,6 +52,12 @@ pub struct Config { pub poll_interval: String, #[serde(default = "default_run_timeout")] pub run_timeout: String, + /// Overrides the default timeout reply text. + #[serde(default)] + pub timeout_reply: Option, + /// Shell command run on run timeout; its stdout becomes the reply. + #[serde(default)] + pub timeout_hook: Option, #[serde(default)] pub self_handles: Vec, #[serde(default)] @@ -903,6 +909,8 @@ mod tests { fn config() -> Config { let root = temp_dir("config-draft-boundary"); Config { + timeout_reply: None, + timeout_hook: None, channel: "imessage".to_string(), channels: Vec::new(), primary_delivery: None, diff --git a/src/gateway/tests.rs b/src/gateway/tests.rs index c7a0bd0..0bbcb3c 100644 --- a/src/gateway/tests.rs +++ b/src/gateway/tests.rs @@ -3956,6 +3956,8 @@ async fn retired_job_approval_reply_explains_direct_creation() { fn test_config(state_path: &str, _sessions_dir: &str, assistant_dir: &str) -> Config { Config { + timeout_reply: None, + timeout_hook: None, channel: "imessage".to_string(), channels: Vec::new(), primary_delivery: None, @@ -4083,6 +4085,135 @@ fn approval_question( .unwrap() } +/// Runs one message against a never-released FakeRunner so the gateway times +/// the run out, then returns the delivered replies and the canonical work dir. +async fn run_to_timeout( + state_tag: &str, + mutate_cfg: &dyn Fn(&mut Config), +) -> (Vec<(String, String)>, String) { + let state_path = temp_path(&format!("{state_tag}-state")); + let state = state_path.to_string_lossy().to_string(); + let assistant_dir = temp_path(&format!("{state_tag}-assistant")); + std::fs::create_dir_all(&assistant_dir).unwrap(); + let mut cfg = test_config(&state, "", &assistant_dir.to_string_lossy()); + mutate_cfg(&mut cfg); + let mut gateway = Gateway::new(cfg).unwrap(); + let mut runners = HashMap::new(); + runners.insert( + AgentBackend::Codex, + Runner::Fake(FakeRunner { + backend: AgentBackend::Codex, + session_id: "fake-session".to_string(), + calls: Arc::new(Mutex::new(Vec::new())), + before_return: None, + wait_for_release: Some(Arc::new(tokio::sync::Notify::new())), + failure: None, + resume_missing_once: None, + }), + ); + gateway.ctx.runners = Arc::new(runners); + gateway.ctx.run_timeout = Duration::from_millis(100); + gateway + .tick_fake(vec![message(1, "me@icloud.com", "", true, "slow")]) + .await; + tokio::time::timeout(Duration::from_secs(2), async { + loop { + if !gateway.ctx.sent_replies.lock().unwrap().is_empty() { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("timeout reply should be delivered"); + let work_dir = std::fs::canonicalize(&assistant_dir) + .unwrap() + .to_string_lossy() + .to_string(); + gateway.queues.clear(); + gateway.drain_workers().await; + let replies = gateway.ctx.sent_replies.lock().unwrap().clone(); + let _ = std::fs::remove_dir_all(&assistant_dir); + let _ = std::fs::remove_dir_all(format!("{state}.cache")); + let _ = std::fs::remove_dir_all(format!("{state}.jobs")); + let _ = std::fs::remove_dir_all(format!("{state}.run")); + for suffix in ["", ".db", ".audit.jsonl", ".audit.jsonl.lock", ".home"] { + let _ = std::fs::remove_file(format!("{state}{suffix}")); + } + (replies, work_dir) +} + +#[tokio::test(flavor = "current_thread")] +async fn timeout_uses_custom_reply() { + let (replies, _) = run_to_timeout("timeout-custom", &|cfg| { + cfg.timeout_reply = Some("custom timeout text".to_string()); + }) + .await; + assert!(replies + .iter() + .any(|(_, reply)| reply.contains("custom timeout text"))); +} + +#[tokio::test(flavor = "current_thread")] +async fn timeout_hook_stdout_wins() { + let cli = crate::test_support::FakeCli::new("hook-wins", "#!/bin/sh\necho hook says hi\n"); + let (replies, _) = run_to_timeout("timeout-hook", &|cfg| { + cfg.timeout_reply = Some("custom timeout text".to_string()); + cfg.timeout_hook = Some(cli.bin()); + }) + .await; + assert!(replies + .iter() + .any(|(_, reply)| reply.contains("hook says hi"))); + assert!(!replies + .iter() + .any(|(_, reply)| reply.contains("custom timeout text"))); +} + +#[tokio::test(flavor = "current_thread")] +async fn timeout_hook_failure_falls_back() { + let cli = crate::test_support::FakeCli::new("hook-fail", "#!/bin/sh\nexit 1\n"); + let (replies, _) = run_to_timeout("timeout-hook-fail", &|cfg| { + cfg.timeout_reply = Some("custom timeout text".to_string()); + cfg.timeout_hook = Some(cli.bin()); + }) + .await; + assert!(replies + .iter() + .any(|(_, reply)| reply.contains("custom timeout text"))); + + let (replies, _) = run_to_timeout("timeout-hook-fail-default", &|cfg| { + cfg.timeout_hook = Some(cli.bin()); + }) + .await; + assert!(replies + .iter() + .any(|(_, reply)| reply.contains("That took too long and was stopped"))); +} + +#[tokio::test(flavor = "current_thread")] +async fn timeout_hook_receives_env_vars() { + let env_path = temp_path("hook-env"); + let script = format!( + "#!/bin/sh\nenv | grep '^PUSH_' > '{}'\necho ok\n", + env_path.to_string_lossy() + ); + let cli = crate::test_support::FakeCli::new("hook-env", &script); + let (replies, work_dir) = run_to_timeout("timeout-hook-env", &|cfg| { + cfg.timeout_hook = Some(cli.bin()); + }) + .await; + assert!(replies.iter().any(|(_, reply)| reply.contains("ok"))); + let env = std::fs::read_to_string(&env_path).unwrap(); + let _ = std::fs::remove_file(&env_path); + assert!(env + .lines() + .any(|line| line.starts_with("PUSH_THREAD=imessage:"))); + assert!(env.contains("PUSH_ROW_ID=1")); + assert!(env.contains("PUSH_BACKEND=codex")); + assert!(env.contains(&format!("PUSH_WORK_DIR={work_dir}"))); +} + fn message(row_id: i64, chat: &str, handle: &str, is_from_me: bool, text: &str) -> RawMessage { RawMessage { row_id, diff --git a/src/gateway/worker.rs b/src/gateway/worker.rs index ecba8df..48cd7cd 100644 --- a/src/gateway/worker.rs +++ b/src/gateway/worker.rs @@ -468,11 +468,11 @@ where format!("{} run timed out", runner.label()), ), ); - let reply = "That took too long and was stopped. Try again or simplify the request."; + let reply = timeout_reply(ctx, &job, &work_dir).await; finish_run_with_gateway_reply( ctx, &job, - reply, + &reply, ReplyLabels { record: "record timeout reply", deliver: "deliver timeout reply", @@ -698,6 +698,48 @@ struct ReplyLabels { completion: &'static str, } +const DEFAULT_TIMEOUT_REPLY: &str = + "That took too long and was stopped. Try again or simplify the request."; + +/// Resolves the reply for a timed-out run: hook stdout (trimmed) wins, then +/// `timeout_reply`, then the default. Hook problems never lose the message. +async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { + let mut reply = ctx + .cfg + .timeout_reply + .clone() + .unwrap_or_else(|| DEFAULT_TIMEOUT_REPLY.to_string()); + let Some(hook) = ctx.cfg.timeout_hook.clone() else { + return reply; + }; + let spawned = tokio::process::Command::new(&hook) + .env("PUSH_THREAD", &job.thread) + .env("PUSH_ROW_ID", job.row_id.to_string()) + .env("PUSH_BACKEND", job.backend.as_str()) + .env("PUSH_WORK_DIR", work_dir) + .output(); + match tokio::time::timeout(Duration::from_secs(5), spawned).await { + Ok(Ok(output)) if output.status.success() => { + let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if stdout.is_empty() { + warn!("[{}] timeout hook produced no output", job.thread); + } else { + reply = stdout; + } + } + Ok(Ok(output)) => warn!( + "[{}] timeout hook exited with {}: using fallback reply", + job.thread, output.status + ), + Ok(Err(error)) => warn!("[{}] timeout hook failed to run: {error}", job.thread), + Err(_) => warn!( + "[{}] timeout hook exceeded 5s: using fallback reply", + job.thread + ), + } + reply +} + /// Records a gateway-authored reply, delivers it, and completes the row. /// Shared by the timeout and failure arms of `handle`. async fn finish_run_with_gateway_reply(ctx: &Ctx, job: &Job, reply: &str, labels: ReplyLabels) { diff --git a/src/test_support.rs b/src/test_support.rs index 9d8d601..df28c40 100644 --- a/src/test_support.rs +++ b/src/test_support.rs @@ -43,6 +43,8 @@ impl Drop for FakeCli { /// touch the filesystem through it. pub fn test_config() -> crate::config::Config { crate::config::Config { + timeout_reply: None, + timeout_hook: None, channel: "imessage".to_string(), channels: Vec::new(), primary_delivery: None, From 4655d63b53c8715deb5432c4137545bef05dcc94 Mon Sep 17 00:00:00 2001 From: Bernardo Kuri Date: Tue, 18 Aug 2026 19:40:57 -0600 Subject: [PATCH 2/8] test: cover 5s timeout-hook overrun fallback --- src/gateway/tests.rs | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/src/gateway/tests.rs b/src/gateway/tests.rs index 0bbcb3c..dac66d6 100644 --- a/src/gateway/tests.rs +++ b/src/gateway/tests.rs @@ -4116,7 +4116,7 @@ async fn run_to_timeout( gateway .tick_fake(vec![message(1, "me@icloud.com", "", true, "slow")]) .await; - tokio::time::timeout(Duration::from_secs(2), async { + tokio::time::timeout(Duration::from_secs(10), async { loop { if !gateway.ctx.sent_replies.lock().unwrap().is_empty() { break; @@ -4191,6 +4191,22 @@ async fn timeout_hook_failure_falls_back() { .any(|(_, reply)| reply.contains("That took too long and was stopped"))); } +#[tokio::test(flavor = "current_thread")] +async fn timeout_hook_overrun_falls_back() { + // The hook sleeps past the 5s budget; the reply must still be delivered. + let cli = crate::test_support::FakeCli::new("hook-slow", "#!/bin/sh\nsleep 30\n"); + let started = std::time::Instant::now(); + let (replies, _) = run_to_timeout("timeout-hook-slow", &|cfg| { + cfg.timeout_reply = Some("custom timeout text".to_string()); + cfg.timeout_hook = Some(cli.bin()); + }) + .await; + assert!(started.elapsed() < std::time::Duration::from_secs(15)); + assert!(replies + .iter() + .any(|(_, reply)| reply.contains("custom timeout text"))); +} + #[tokio::test(flavor = "current_thread")] async fn timeout_hook_receives_env_vars() { let env_path = temp_path("hook-env"); From ac7e22399e2b2cf6f086dcdbc3bda421c67c4116 Mon Sep 17 00:00:00 2001 From: Bernardo Kuri Date: Thu, 20 Aug 2026 13:59:35 -0600 Subject: [PATCH 3/8] fix: review findings on the timeout hook - Run timeout_hook through /bin/sh -c so documented shell syntax (arguments, pipes, expansion) works, not just bare paths - kill_on_drop + dedicated process group so an over-budget hook is stopped when the 5s future is dropped - Treat a blank timeout_reply as unset instead of delivering an empty message - Pin the overrun test to the 5s budget (4.5-8s window) --- docs/configuration.md | 2 +- src/gateway/tests.rs | 33 +++++++++++++++++++++++++++++++-- src/gateway/worker.rs | 16 ++++++++++++++-- 3 files changed, 46 insertions(+), 5 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index fcdc233..9f64d1c 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -210,7 +210,7 @@ requests. Review [permissions and security](security.md) before enabling jobs. | `poll_interval` | `"3s"` | Delay between channel polls | | `run_timeout` | `"10m"` | Maximum chat backend run time | | `timeout_reply` | hardcoded string | Overrides the default timeout reply text | -| `timeout_hook` | none | Shell command run on run timeout; stdout becomes the reply (env vars: `PUSH_THREAD`, `PUSH_ROW_ID`, `PUSH_BACKEND`, `PUSH_WORK_DIR`; 5s budget, warn and fall back on any failure) | +| `timeout_hook` | none | Shell command run via `/bin/sh -c` on run timeout; stdout becomes the reply (env vars: `PUSH_THREAD`, `PUSH_ROW_ID`, `PUSH_BACKEND`, `PUSH_WORK_DIR`; 5s budget, warn and fall back on any failure) | ### iMessage diff --git a/src/gateway/tests.rs b/src/gateway/tests.rs index dac66d6..65ccbf1 100644 --- a/src/gateway/tests.rs +++ b/src/gateway/tests.rs @@ -4193,7 +4193,8 @@ async fn timeout_hook_failure_falls_back() { #[tokio::test(flavor = "current_thread")] async fn timeout_hook_overrun_falls_back() { - // The hook sleeps past the 5s budget; the reply must still be delivered. + // The hook sleeps past the 5s budget; the reply must still be delivered + // close to the budget itself (5s plus CI margin, not 15s). let cli = crate::test_support::FakeCli::new("hook-slow", "#!/bin/sh\nsleep 30\n"); let started = std::time::Instant::now(); let (replies, _) = run_to_timeout("timeout-hook-slow", &|cfg| { @@ -4201,12 +4202,40 @@ async fn timeout_hook_overrun_falls_back() { cfg.timeout_hook = Some(cli.bin()); }) .await; - assert!(started.elapsed() < std::time::Duration::from_secs(15)); + let elapsed = started.elapsed(); + assert!(elapsed >= std::time::Duration::from_millis(4_500)); + assert!(elapsed < std::time::Duration::from_secs(8)); assert!(replies .iter() .any(|(_, reply)| reply.contains("custom timeout text"))); } +#[tokio::test(flavor = "current_thread")] +async fn timeout_hook_accepts_shell_syntax() { + // The hook value runs through /bin/sh -c, so arguments and expansion + // work without a wrapper script. + let (replies, _) = run_to_timeout("timeout-hook-shell", &|cfg| { + cfg.timeout_hook = Some("echo shell-hook-$PUSH_BACKEND".to_string()); + }) + .await; + assert!(replies + .iter() + .any(|(_, reply)| reply.contains("shell-hook-codex"))); +} + +#[tokio::test(flavor = "current_thread")] +async fn blank_timeout_reply_falls_back_to_default() { + // A whitespace-only timeout_reply is treated as unset rather than + // delivering an empty message. + let (replies, _) = run_to_timeout("timeout-blank-reply", &|cfg| { + cfg.timeout_reply = Some(" ".to_string()); + }) + .await; + assert!(replies + .iter() + .any(|(_, reply)| reply.contains("That took too long and was stopped"))); +} + #[tokio::test(flavor = "current_thread")] async fn timeout_hook_receives_env_vars() { let env_path = temp_path("hook-env"); diff --git a/src/gateway/worker.rs b/src/gateway/worker.rs index 48cd7cd..3245a9b 100644 --- a/src/gateway/worker.rs +++ b/src/gateway/worker.rs @@ -707,16 +707,28 @@ async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { let mut reply = ctx .cfg .timeout_reply - .clone() + .as_deref() + .map(str::trim) + .filter(|reply| !reply.is_empty()) + .map(str::to_string) .unwrap_or_else(|| DEFAULT_TIMEOUT_REPLY.to_string()); let Some(hook) = ctx.cfg.timeout_hook.clone() else { return reply; }; - let spawned = tokio::process::Command::new(&hook) + // The hook is documented as a shell command: run it through /bin/sh so + // arguments, pipes, and redirects work. kill_on_drop + its own process + // group stop an over-budget hook when the timeout drops the future. + // ponytail: kills the direct child only; a group-wide kill(-pgid) if a + // hook reliably leaves grandchildren behind. + let spawned = tokio::process::Command::new("/bin/sh") + .arg("-c") + .arg(&hook) .env("PUSH_THREAD", &job.thread) .env("PUSH_ROW_ID", job.row_id.to_string()) .env("PUSH_BACKEND", job.backend.as_str()) .env("PUSH_WORK_DIR", work_dir) + .kill_on_drop(true) + .process_group(0) .output(); match tokio::time::timeout(Duration::from_secs(5), spawned).await { Ok(Ok(output)) if output.status.success() => { From 825e438225c4068d5a34b0344fa3e9ac13ebf595 Mon Sep 17 00:00:00 2001 From: Bernardo Kuri Date: Thu, 20 Aug 2026 14:09:43 -0600 Subject: [PATCH 4/8] fix: bound captured timeout-hook output Route hook stdout through head -c 65536 with pipefail so a runaway hook SIGPIPEs instead of exhausting gateway memory; overflow lands in the existing warn-and-fallback path. Hook stderr is discarded. --- src/gateway/tests.rs | 18 ++++++++++++++++++ src/gateway/worker.rs | 10 +++++++++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/src/gateway/tests.rs b/src/gateway/tests.rs index 65ccbf1..79edd7d 100644 --- a/src/gateway/tests.rs +++ b/src/gateway/tests.rs @@ -4223,6 +4223,24 @@ async fn timeout_hook_accepts_shell_syntax() { .any(|(_, reply)| reply.contains("shell-hook-codex"))); } +#[tokio::test(flavor = "current_thread")] +async fn timeout_hook_unbounded_output_falls_back() { + // A hook that spews without end must not exhaust memory or deliver a + // giant reply: head -c caps capture, pipefail fails the pipeline, the + // fallback reply is used. + let (replies, _) = run_to_timeout("timeout-hook-runaway", &|cfg| { + cfg.timeout_reply = Some("custom timeout text".to_string()); + cfg.timeout_hook = Some("yes runaway".to_string()); + }) + .await; + assert!(replies + .iter() + .any(|(_, reply)| reply.contains("custom timeout text"))); + assert!(replies + .iter() + .all(|(_, reply)| !reply.contains("runawayrunaway"))); +} + #[tokio::test(flavor = "current_thread")] async fn blank_timeout_reply_falls_back_to_default() { // A whitespace-only timeout_reply is treated as unset rather than diff --git a/src/gateway/worker.rs b/src/gateway/worker.rs index 3245a9b..059f953 100644 --- a/src/gateway/worker.rs +++ b/src/gateway/worker.rs @@ -720,13 +720,21 @@ async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { // group stop an over-budget hook when the timeout drops the future. // ponytail: kills the direct child only; a group-wide kill(-pgid) if a // hook reliably leaves grandchildren behind. + // The hook is documented as a shell command: run it through /bin/sh so + // arguments, pipes, and redirects work. `head -c` bounds captured stdout + // (a runaway hook SIGPIPEs instead of exhausting memory; pipefail then + // routes it to the fallback). kill_on_drop + its own process group stop + // an over-budget hook when the timeout drops the future. + // ponytail: kills the direct child only; a group-wide kill(-pgid) if a + // hook reliably leaves grandchildren behind. let spawned = tokio::process::Command::new("/bin/sh") .arg("-c") - .arg(&hook) + .arg(format!("set -o pipefail; {hook} | head -c 65536")) .env("PUSH_THREAD", &job.thread) .env("PUSH_ROW_ID", job.row_id.to_string()) .env("PUSH_BACKEND", job.backend.as_str()) .env("PUSH_WORK_DIR", work_dir) + .stderr(std::process::Stdio::null()) .kill_on_drop(true) .process_group(0) .output(); From 99a5ba90799d13d08d527c7158969dff360a4970 Mon Sep 17 00:00:00 2001 From: Bernardo Kuri Date: Thu, 20 Aug 2026 14:16:27 -0600 Subject: [PATCH 5/8] fix: bound hook stdout in Rust, not via pipefail set -o pipefail is not POSIX and dash < 0.5.11 (Ubuntu 20.04, older Debian) rejects it, silently disabling every hook on those systems. Read hook stdout with a 64KiB byte cap instead: exceeding the cap kills the hook and uses the fallback reply. Same memory bound, works on any POSIX sh. --- src/gateway/worker.rs | 72 +++++++++++++++++++++++++++++++++++++------ 1 file changed, 62 insertions(+), 10 deletions(-) diff --git a/src/gateway/worker.rs b/src/gateway/worker.rs index 059f953..a6099b0 100644 --- a/src/gateway/worker.rs +++ b/src/gateway/worker.rs @@ -701,6 +701,8 @@ struct ReplyLabels { const DEFAULT_TIMEOUT_REPLY: &str = "That took too long and was stopped. Try again or simplify the request."; +const MAX_HOOK_STDOUT: usize = 64 * 1024; + /// Resolves the reply for a timed-out run: hook stdout (trimmed) wins, then /// `timeout_reply`, then the default. Hook problems never lose the message. async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { @@ -727,31 +729,81 @@ async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { // an over-budget hook when the timeout drops the future. // ponytail: kills the direct child only; a group-wide kill(-pgid) if a // hook reliably leaves grandchildren behind. - let spawned = tokio::process::Command::new("/bin/sh") + // The hook is documented as a shell command: run it through /bin/sh so + // arguments, pipes, and redirects work (POSIX sh only — no pipefail, + // which older dash rejects). stdout is read with a byte cap so a runaway + // hook cannot exhaust memory; exceeding the cap kills the hook and uses + // the fallback reply. kill_on_drop + its own process group stop an + // over-budget hook when the timeout drops the future. + // ponytail: kills the direct child only; a group-wide kill(-pgid) if a + // hook reliably leaves grandchildren behind. + let child = tokio::process::Command::new("/bin/sh") .arg("-c") - .arg(format!("set -o pipefail; {hook} | head -c 65536")) + .arg(&hook) .env("PUSH_THREAD", &job.thread) .env("PUSH_ROW_ID", job.row_id.to_string()) .env("PUSH_BACKEND", job.backend.as_str()) .env("PUSH_WORK_DIR", work_dir) + .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::null()) .kill_on_drop(true) .process_group(0) - .output(); - match tokio::time::timeout(Duration::from_secs(5), spawned).await { - Ok(Ok(output)) if output.status.success() => { - let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string(); + .spawn(); + let mut child = match child { + Ok(child) => child, + Err(error) => { + warn!("[{}] timeout hook failed to run: {error}", job.thread); + return reply; + } + }; + let mut stdout_pipe = child.stdout.take().expect("hook stdout piped"); + let collect = async { + use tokio::io::AsyncReadExt; + let mut stdout = Vec::new(); + let mut chunk = [0u8; 8192]; + let mut capped = false; + loop { + match stdout_pipe.read(&mut chunk).await { + Ok(0) | Err(_) => break, + Ok(n) => { + if stdout.len() + n > MAX_HOOK_STDOUT { + capped = true; + break; + } + stdout.extend_from_slice(&chunk[..n]); + } + } + } + (stdout, capped) + }; + let result = tokio::time::timeout(Duration::from_secs(5), async { + let (stdout, capped) = collect.await; + if capped { + let _ = child.kill().await; + } + (stdout, capped, child.wait().await) + }) + .await; + match result { + Ok((stdout, false, Ok(status))) if status.success() => { + let stdout = String::from_utf8_lossy(&stdout).trim().to_string(); if stdout.is_empty() { warn!("[{}] timeout hook produced no output", job.thread); } else { reply = stdout; } } - Ok(Ok(output)) => warn!( - "[{}] timeout hook exited with {}: using fallback reply", - job.thread, output.status + Ok((_, true, _)) => warn!( + "[{}] timeout hook exceeded {} bytes: using fallback reply", + job.thread, MAX_HOOK_STDOUT + ), + Ok((_, _, Ok(status))) => warn!( + "[{}] timeout hook exited with {status}: using fallback reply", + job.thread ), - Ok(Err(error)) => warn!("[{}] timeout hook failed to run: {error}", job.thread), + Ok((_, _, Err(error))) => { + warn!("[{}] timeout hook failed to run: {error}", job.thread) + } Err(_) => warn!( "[{}] timeout hook exceeded 5s: using fallback reply", job.thread From ac82729522e0d7da90f70293bdce4be1b562915b Mon Sep 17 00:00:00 2001 From: Bernardo Kuri <138886+bkuri@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:51:14 -0600 Subject: [PATCH 6/8] Clean up comments in worker.rs Removed redundant comments about hook execution and memory management. --- src/gateway/worker.rs | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/src/gateway/worker.rs b/src/gateway/worker.rs index a6099b0..59a196e 100644 --- a/src/gateway/worker.rs +++ b/src/gateway/worker.rs @@ -718,24 +718,12 @@ async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { return reply; }; // The hook is documented as a shell command: run it through /bin/sh so - // arguments, pipes, and redirects work. kill_on_drop + its own process - // group stop an over-budget hook when the timeout drops the future. - // ponytail: kills the direct child only; a group-wide kill(-pgid) if a - // hook reliably leaves grandchildren behind. - // The hook is documented as a shell command: run it through /bin/sh so - // arguments, pipes, and redirects work. `head -c` bounds captured stdout - // (a runaway hook SIGPIPEs instead of exhausting memory; pipefail then - // routes it to the fallback). kill_on_drop + its own process group stop - // an over-budget hook when the timeout drops the future. - // ponytail: kills the direct child only; a group-wide kill(-pgid) if a - // hook reliably leaves grandchildren behind. - // The hook is documented as a shell command: run it through /bin/sh so // arguments, pipes, and redirects work (POSIX sh only — no pipefail, // which older dash rejects). stdout is read with a byte cap so a runaway // hook cannot exhaust memory; exceeding the cap kills the hook and uses // the fallback reply. kill_on_drop + its own process group stop an // over-budget hook when the timeout drops the future. - // ponytail: kills the direct child only; a group-wide kill(-pgid) if a + // note: kills the direct child only; a group-wide kill(-pgid) if a // hook reliably leaves grandchildren behind. let child = tokio::process::Command::new("/bin/sh") .arg("-c") From d8dbc7d95c73e553d97813dbc402fd5d63f3c464 Mon Sep 17 00:00:00 2001 From: Bernardo Kuri Date: Thu, 20 Aug 2026 20:59:03 -0600 Subject: [PATCH 7/8] fix: treat hook stdout read error as hook failure A read error on the hook's stdout pipe was swallowed as EOF, so a hook that wrote partial output before the pipe broke could deliver those partial bytes as a successful reply. read_hook_stdout now returns the first read error; the worker logs it and uses the fallback reply. --- src/gateway/tests.rs | 50 ++++++++++++++++++++++++++++++++ src/gateway/worker.rs | 67 +++++++++++++++++++++++++++---------------- 2 files changed, 92 insertions(+), 25 deletions(-) diff --git a/src/gateway/tests.rs b/src/gateway/tests.rs index 79edd7d..2a0a280 100644 --- a/src/gateway/tests.rs +++ b/src/gateway/tests.rs @@ -4241,6 +4241,56 @@ async fn timeout_hook_unbounded_output_falls_back() { .all(|(_, reply)| !reply.contains("runawayrunaway"))); } +/// Reader that yields bytes once, then fails — models a hook pipe that +/// breaks mid-write instead of reaching EOF. +struct FailingPipeReader { + bytes: Vec, + done: bool, +} + +impl tokio::io::AsyncRead for FailingPipeReader { + fn poll_read( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + buf: &mut tokio::io::ReadBuf<'_>, + ) -> std::task::Poll> { + let this = self.get_mut(); + if this.done { + return std::task::Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "pipe broke mid-write", + ))); + } + this.done = true; + buf.put_slice(&this.bytes); + std::task::Poll::Ready(Ok(())) + } +} + +#[tokio::test(flavor = "current_thread")] +async fn hook_stdout_read_error_surfaces_not_swallowed() { + // A read error after partial output must be returned to the caller so + // the worker logs it and uses the fallback reply; treating it as EOF + // would deliver the partial bytes as a successful hook result. + use super::worker::read_hook_stdout; + + let mut reader = FailingPipeReader { + bytes: b"partial".to_vec(), + done: false, + }; + let (stdout, capped, read_error) = read_hook_stdout(&mut reader).await; + assert_eq!(stdout, b"partial"); + assert!(!capped); + assert!(read_error.is_some()); + + // Clean EOF stays clean: no error, all bytes. + let mut eof = std::io::Cursor::new(b"all good".to_vec()); + let (stdout, capped, read_error) = read_hook_stdout(&mut eof).await; + assert_eq!(stdout, b"all good"); + assert!(!capped); + assert!(read_error.is_none()); +} + #[tokio::test(flavor = "current_thread")] async fn blank_timeout_reply_falls_back_to_default() { // A whitespace-only timeout_reply is treated as unset rather than diff --git a/src/gateway/worker.rs b/src/gateway/worker.rs index 59a196e..15736a5 100644 --- a/src/gateway/worker.rs +++ b/src/gateway/worker.rs @@ -703,6 +703,37 @@ const DEFAULT_TIMEOUT_REPLY: &str = const MAX_HOOK_STDOUT: usize = 64 * 1024; +/// Reads hook stdout with a byte cap, returning the bytes, whether the cap +/// was exceeded, and the first read error (if any). A read error is not +/// EOF: partial output from a failing pipe is a hook failure, so the error +/// is surfaced for the caller to log and fall back. +pub(crate) async fn read_hook_stdout( + reader: &mut R, +) -> (Vec, bool, Option) { + use tokio::io::AsyncReadExt; + let mut stdout = Vec::new(); + let mut chunk = [0u8; 8192]; + let mut capped = false; + let mut read_error = None; + loop { + match reader.read(&mut chunk).await { + Ok(0) => break, + Ok(n) => { + if stdout.len() + n > MAX_HOOK_STDOUT { + capped = true; + break; + } + stdout.extend_from_slice(&chunk[..n]); + } + Err(error) => { + read_error = Some(error); + break; + } + } + } + (stdout, capped, read_error) +} + /// Resolves the reply for a timed-out run: hook stdout (trimmed) wins, then /// `timeout_reply`, then the default. Hook problems never lose the message. async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { @@ -745,35 +776,21 @@ async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { } }; let mut stdout_pipe = child.stdout.take().expect("hook stdout piped"); - let collect = async { - use tokio::io::AsyncReadExt; - let mut stdout = Vec::new(); - let mut chunk = [0u8; 8192]; - let mut capped = false; - loop { - match stdout_pipe.read(&mut chunk).await { - Ok(0) | Err(_) => break, - Ok(n) => { - if stdout.len() + n > MAX_HOOK_STDOUT { - capped = true; - break; - } - stdout.extend_from_slice(&chunk[..n]); - } - } - } - (stdout, capped) - }; + let collect = read_hook_stdout(&mut stdout_pipe); let result = tokio::time::timeout(Duration::from_secs(5), async { - let (stdout, capped) = collect.await; + let (stdout, capped, read_error) = collect.await; if capped { let _ = child.kill().await; } - (stdout, capped, child.wait().await) + (stdout, capped, read_error, child.wait().await) }) .await; match result { - Ok((stdout, false, Ok(status))) if status.success() => { + Ok((_, _, Some(read_error), _)) => warn!( + "[{}] timeout hook stdout read failed: {read_error}: using fallback reply", + job.thread + ), + Ok((stdout, false, None, Ok(status))) if status.success() => { let stdout = String::from_utf8_lossy(&stdout).trim().to_string(); if stdout.is_empty() { warn!("[{}] timeout hook produced no output", job.thread); @@ -781,15 +798,15 @@ async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { reply = stdout; } } - Ok((_, true, _)) => warn!( + Ok((_, true, _, _)) => warn!( "[{}] timeout hook exceeded {} bytes: using fallback reply", job.thread, MAX_HOOK_STDOUT ), - Ok((_, _, Ok(status))) => warn!( + Ok((_, _, _, Ok(status))) => warn!( "[{}] timeout hook exited with {status}: using fallback reply", job.thread ), - Ok((_, _, Err(error))) => { + Ok((_, _, _, Err(error))) => { warn!("[{}] timeout hook failed to run: {error}", job.thread) } Err(_) => warn!( From 847c4b4f6832eeb82ab061981a214865d7784129 Mon Sep 17 00:00:00 2001 From: Bernardo Kuri Date: Thu, 20 Aug 2026 21:48:44 -0600 Subject: [PATCH 8/8] fix: kill the hook process group, not just the shell leader kill_on_drop and child.kill() only signal the direct child, so a hook that backgrounded work (e.g. 'sleep 60 &') leaked descendants past the 5s budget. The hook already runs in its own process group (process_group(0), pgid == leader pid); every failure path now sends SIGKILL to -pgid: the output-cap branch and the 5s-timeout expiry (where kill_on_drop reaped the leader but not the group). Regression test starts a background descendant and asserts it dies. --- src/gateway/tests.rs | 39 +++++++++++++++++++++++++++++++++++++++ src/gateway/worker.rs | 38 +++++++++++++++++++++++++++++--------- 2 files changed, 68 insertions(+), 9 deletions(-) diff --git a/src/gateway/tests.rs b/src/gateway/tests.rs index 2a0a280..8f60c50 100644 --- a/src/gateway/tests.rs +++ b/src/gateway/tests.rs @@ -4210,6 +4210,45 @@ async fn timeout_hook_overrun_falls_back() { .any(|(_, reply)| reply.contains("custom timeout text"))); } +#[tokio::test(flavor = "current_thread")] +async fn timeout_hook_descendants_are_killed() { + // A hook that backgrounds a long sleep must not leak it past the 5s + // budget: the whole process group is signalled, not just the shell. + // (pid file outside $PUSH_WORK_DIR: run_to_timeout wipes the work dir.) + let pid_path = temp_path("timeout-hook-descendant-pid"); + let hook = format!( + "sleep 60 & echo $! > '{}'; sleep 30", + pid_path.to_string_lossy() + ); + let (replies, _) = run_to_timeout("timeout-hook-descendant", &|cfg| { + cfg.timeout_reply = Some("custom timeout text".to_string()); + cfg.timeout_hook = Some(hook.clone()); + }) + .await; + assert!(replies + .iter() + .any(|(_, reply)| reply.contains("custom timeout text"))); + + let pid: libc::pid_t = std::fs::read_to_string(&pid_path) + .unwrap() + .trim() + .parse() + .unwrap(); + let _ = std::fs::remove_file(&pid_path); + // SIGKILL delivery is asynchronous; poll briefly for the descendant to + // disappear (signal 0 probes existence without sending anything). + let mut dead = false; + for _ in 0..100 { + // Safety: signal syscall with an integer pid; no pointers involved. + if unsafe { libc::kill(pid, 0) } != 0 { + dead = true; + break; + } + std::thread::sleep(std::time::Duration::from_millis(20)); + } + assert!(dead, "hook descendant {pid} survived the timeout kill"); +} + #[tokio::test(flavor = "current_thread")] async fn timeout_hook_accepts_shell_syntax() { // The hook value runs through /bin/sh -c, so arguments and expansion diff --git a/src/gateway/worker.rs b/src/gateway/worker.rs index 15736a5..6bcc0f6 100644 --- a/src/gateway/worker.rs +++ b/src/gateway/worker.rs @@ -703,6 +703,17 @@ const DEFAULT_TIMEOUT_REPLY: &str = const MAX_HOOK_STDOUT: usize = 64 * 1024; +/// SIGKILLs the hook's whole process group (negative pid). Covers the shell +/// leader and any backgrounded descendants — child.kill()/kill_on_drop only +/// ever signal the leader. Errors are ignored: the group may already be gone +/// on a given path. +fn kill_hook_group(pgid: libc::pid_t) { + // Safety: signal syscall with an integer pid; no pointers involved. + unsafe { + libc::kill(-pgid, libc::SIGKILL); + } +} + /// Reads hook stdout with a byte cap, returning the bytes, whether the cap /// was exceeded, and the first read error (if any). A read error is not /// EOF: partial output from a failing pipe is a hook failure, so the error @@ -752,10 +763,10 @@ async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { // arguments, pipes, and redirects work (POSIX sh only — no pipefail, // which older dash rejects). stdout is read with a byte cap so a runaway // hook cannot exhaust memory; exceeding the cap kills the hook and uses - // the fallback reply. kill_on_drop + its own process group stop an - // over-budget hook when the timeout drops the future. - // note: kills the direct child only; a group-wide kill(-pgid) if a - // hook reliably leaves grandchildren behind. + // the fallback reply. The hook runs in its own process group + // (process_group(0), leader pid == pgid) and every failure path + // signals the whole group, so backgrounded descendants die with the + // shell instead of leaking past the budget. let child = tokio::process::Command::new("/bin/sh") .arg("-c") .arg(&hook) @@ -776,11 +787,15 @@ async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { } }; let mut stdout_pipe = child.stdout.take().expect("hook stdout piped"); + // Leader pid == process group id (process_group(0)); kept outside the + // timeout scope so the expiry path can still signal the group after the + // future (and the Child) is dropped. + let hook_pgid = child.id().expect("hook child alive before wait") as libc::pid_t; let collect = read_hook_stdout(&mut stdout_pipe); let result = tokio::time::timeout(Duration::from_secs(5), async { let (stdout, capped, read_error) = collect.await; if capped { - let _ = child.kill().await; + kill_hook_group(hook_pgid); } (stdout, capped, read_error, child.wait().await) }) @@ -809,10 +824,15 @@ async fn timeout_reply(ctx: &Ctx, job: &Job, work_dir: &str) -> String { Ok((_, _, _, Err(error))) => { warn!("[{}] timeout hook failed to run: {error}", job.thread) } - Err(_) => warn!( - "[{}] timeout hook exceeded 5s: using fallback reply", - job.thread - ), + Err(_) => { + // kill_on_drop reaped the leader when the inner future dropped; + // the group signal takes care of any descendants it spawned. + kill_hook_group(hook_pgid); + warn!( + "[{}] timeout hook exceeded 5s: using fallback reply", + job.thread + ) + } } reply }