From 9a8100cb4696c51e50097c9a8cf4c7950dcc2b37 Mon Sep 17 00:00:00 2001 From: "pullfrog[bot]" <226033991+pullfrog[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 21:37:16 +0000 Subject: [PATCH] =?UTF-8?q?fix(core):=20P0a=20trust=20correctness=20?= =?UTF-8?q?=E2=80=94=20H13=20cancel,=20H5=20write=20content,=20H3=20stream?= =?UTF-8?q?=20caps?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Close Agent OS roadmap P0a items from #8: - H13: cancelled subagent/chain abort paths return non-ok; never promote worktrees unless outcome.ok and the cancel token is clear (single + parallel). - H5: write_file and bulk_write require a present string `content` key (explicit empty string still allowed); missing key leaves existing files. - H3: bound streamed tool-call index (<64) and total args bytes (8 MiB) on OpenAI, Anthropic, Codex, and Gemini stream paths; do not retry cap overflows. --- core/src/provider.rs | 145 ++++++++++++++++++++++++-------- core/src/providers/streaming.rs | 70 +++++++++++++++ core/src/subagent.rs | 41 +++++++-- core/src/tools.rs | 85 ++++++++++++++++++- 4 files changed, 296 insertions(+), 45 deletions(-) diff --git a/core/src/provider.rs b/core/src/provider.rs index ca91a4a..3fc5ab8 100644 --- a/core/src/provider.rs +++ b/core/src/provider.rs @@ -17,7 +17,9 @@ use crate::providers::adapter::{ pub use crate::providers::discovery::*; use crate::providers::registry::{adapter_for, protocol_for}; use crate::providers::sse::{SseDecoder, SseFrame}; -use crate::providers::streaming::NormalizedStreamEvent; +use crate::providers::streaming::{ + append_stream_tool_args, ensure_stream_tool_slot, NormalizedStreamEvent, +}; pub use crate::providers::usage::*; use futures_util::StreamExt; use serde_json::{json, Value}; @@ -1282,11 +1284,21 @@ fn apply_codex_event( } NormalizedStreamEvent::ToolCallStart(delta) => { timer.mark_first_token(); + if calls.len() >= crate::providers::streaming::MAX_STREAM_TOOL_CALL_INDEX { + return Err(format!( + "stream tool_call count exceeds cap ({})", + crate::providers::streaming::MAX_STREAM_TOOL_CALL_INDEX + )); + } let index = calls.len(); + let args = delta.arguments.unwrap_or_else(|| "{}".into()); + let mut total_args: usize = calls.iter().map(|c| c.args.len()).sum(); + let mut acc_args = String::new(); + append_stream_tool_args(&mut total_args, &mut acc_args, &args)?; let call = ToolAccum { id: delta.id.unwrap_or_default(), name: delta.name.unwrap_or_default(), - args: delta.arguments.unwrap_or_else(|| "{}".into()), + args: acc_args, }; if !quiet { emit( @@ -1408,6 +1420,7 @@ async fn stream_turn_openai( // advanced when minimax_cumulative; otherwise unused. let mut wire_text = String::new(); let mut tool_calls: Vec = Vec::new(); + let mut tool_args_bytes: usize = 0; let mut finish_reason = String::new(); let mut tokens_in: u64 = 0; let mut tokens_out: u64 = 0; @@ -1514,8 +1527,7 @@ async fn stream_turn_openai( if !quiet { emitted = true; emit( - &Event::new("delta") - .with("text", json!(delta)), + &Event::new("delta").with("text", json!(delta)), ); } } @@ -1556,8 +1568,11 @@ async fn stream_turn_openai( | NormalizedStreamEvent::ToolCallDelta(delta) => { timer.mark_first_token(); let idx = delta.index; - while tool_calls.len() <= idx { - tool_calls.push(ToolAccum::default()); + // H3: bound index growth + total args bytes so a + // malicious high index cannot allocate multi-GB. + if let Err(e) = ensure_stream_tool_slot(&mut tool_calls, idx) { + err = Some(e); + break 'read_stream; } let acc = &mut tool_calls[idx]; if let Some(id) = delta.id { @@ -1587,7 +1602,14 @@ async fn stream_turn_openai( } } if let Some(arguments) = delta.arguments { - acc.args.push_str(&arguments); + if let Err(e) = append_stream_tool_args( + &mut tool_args_bytes, + &mut acc.args, + &arguments, + ) { + err = Some(e); + break 'read_stream; + } if !quiet { emitted = true; emit( @@ -1734,9 +1756,10 @@ async fn stream_turn_openai( add_structured_tool_call_recovery_instruction(&mut body); content.clear(); wire_text.clear(); - think_demux = ThinkTagDemux::default(); + think_demux = ThinkTagDemux::default(); reasoning.clear(); tool_calls.clear(); + tool_args_bytes = 0; finish_reason.clear(); tokens_in = 0; tokens_out = 0; @@ -1768,6 +1791,7 @@ async fn stream_turn_openai( wire_text.clear(); think_demux = ThinkTagDemux::default(); tool_calls.clear(); + tool_args_bytes = 0; finish_reason.clear(); tokens_in = 0; tokens_out = 0; @@ -1959,7 +1983,8 @@ async fn stream_turn_gemini( terminal_event |= is_terminal_stream_event(&event); match event { NormalizedStreamEvent::TextDelta(text) => { - if let Some(delta) = append_stream_fragment(&mut content, &text, false) { + if let Some(delta) = append_stream_fragment(&mut content, &text, false) + { if content.len() == delta.len() { timer.mark_first_token(); } @@ -1970,7 +1995,9 @@ async fn stream_turn_gemini( } } NormalizedStreamEvent::ReasoningDelta(text) => { - if let Some(delta) = append_stream_fragment(&mut reasoning, &text, false) { + if let Some(delta) = + append_stream_fragment(&mut reasoning, &text, false) + { if reasoning.len() == delta.len() { timer.mark_first_token(); } @@ -1983,11 +2010,29 @@ async fn stream_turn_gemini( NormalizedStreamEvent::ToolCallStart(delta) => { timer.mark_first_token(); let name = delta.name.unwrap_or_default(); - let args = delta - .arguments - .as_deref() - .and_then(|arguments| serde_json::from_str(arguments).ok()) - .unwrap_or_else(|| json!({})); + let args_raw = delta.arguments.unwrap_or_else(|| "{}".into()); + if genai_tool_calls.len() + >= crate::providers::streaming::MAX_STREAM_TOOL_CALL_INDEX + { + err = Some(format!( + "stream tool_call count exceeds cap ({})", + crate::providers::streaming::MAX_STREAM_TOOL_CALL_INDEX + )); + break 'read_stream; + } + let mut total_args: usize = genai_tool_calls + .iter() + .map(|(_, a)| a.to_string().len()) + .sum(); + let mut acc = String::new(); + if let Err(e) = + append_stream_tool_args(&mut total_args, &mut acc, &args_raw) + { + err = Some(e); + break 'read_stream; + } + let args: Value = + serde_json::from_str(&acc).unwrap_or_else(|_| json!({})); let index = genai_tool_calls.len(); genai_tool_calls.push((name.clone(), args.clone())); if !quiet { @@ -2665,6 +2710,10 @@ fn should_retry_stream_attempt( if attempt >= max_attempts { return false; } + // Policy/DoS caps (H3) are permanent for this request shape — never retry. + if message.contains("exceeds cap") || message.contains("exceed cap") { + return false; + } if !emitted { return true; } @@ -3397,6 +3446,7 @@ async fn stream_turn_anthropic( let mut content = String::new(); let mut reasoning = String::new(); let mut blocks: Vec = Vec::new(); + let mut tool_args_bytes: usize = 0; let mut finish_reason = String::new(); let mut tokens_in: u64 = 0; let mut tokens_out: u64 = 0; @@ -3461,7 +3511,9 @@ async fn stream_turn_anthropic( terminal_event |= is_terminal_stream_event(&event); match event { NormalizedStreamEvent::TextDelta(text) => { - if let Some(delta) = append_stream_fragment(&mut content, &text, minimax) { + if let Some(delta) = + append_stream_fragment(&mut content, &text, minimax) + { if content.len() == delta.len() { timer.mark_first_token(); } @@ -3472,7 +3524,9 @@ async fn stream_turn_anthropic( } } NormalizedStreamEvent::ReasoningDelta(text) => { - if let Some(delta) = append_stream_fragment(&mut reasoning, &text, minimax) { + if let Some(delta) = + append_stream_fragment(&mut reasoning, &text, minimax) + { if reasoning.len() == delta.len() { timer.mark_first_token(); } @@ -3484,8 +3538,9 @@ async fn stream_turn_anthropic( } NormalizedStreamEvent::ToolCallStart(delta) => { timer.mark_first_token(); - while blocks.len() <= delta.index { - blocks.push(AnthropicBlock::default()); + if let Err(e) = ensure_stream_tool_slot(&mut blocks, delta.index) { + err = Some(e); + break 'read_stream; } let block = &mut blocks[delta.index]; block.kind = "tool_use".into(); @@ -3513,13 +3568,21 @@ async fn stream_turn_anthropic( } NormalizedStreamEvent::ToolCallDelta(delta) => { timer.mark_first_token(); - while blocks.len() <= delta.index { - blocks.push(AnthropicBlock::default()); + if let Err(e) = ensure_stream_tool_slot(&mut blocks, delta.index) { + err = Some(e); + break 'read_stream; } let block = &mut blocks[delta.index]; block.kind = "tool_use".into(); if let Some(arguments) = delta.arguments { - block.tool_args.push_str(&arguments); + if let Err(e) = append_stream_tool_args( + &mut tool_args_bytes, + &mut block.tool_args, + &arguments, + ) { + err = Some(e); + break 'read_stream; + } if !quiet { emitted = true; emit( @@ -3600,6 +3663,7 @@ async fn stream_turn_anthropic( content.clear(); reasoning.clear(); blocks.clear(); + tool_args_bytes = 0; finish_reason.clear(); tokens_in = 0; tokens_out = 0; @@ -3688,6 +3752,20 @@ mod tests { assert!(!reasoning_contains_dsml_tool_call("ordinary reasoning")); } + #[test] + fn stream_tool_slot_helpers_reject_unbounded_index() { + // H3: high index must fail without growing the vector to multi-GB. + let mut calls: Vec = Vec::new(); + assert!(ensure_stream_tool_slot(&mut calls, 0).is_ok()); + let err = ensure_stream_tool_slot( + &mut calls, + crate::providers::streaming::MAX_STREAM_TOOL_CALL_INDEX, + ) + .unwrap_err(); + assert!(err.contains("exceeds cap"), "{err}"); + assert_eq!(calls.len(), 1); + } + #[test] fn dsml_recovery_rejects_unknown_tools_and_invalid_json() { let registered = vec![json!({"function": {"name": "edit"}})]; @@ -4560,34 +4638,24 @@ mod tests { assert!(!is_minimax("https://api.openai.com/v1")); } - #[test] fn think_tag_demux_splits_complete_and_partial_tags() { let mut d = ThinkTagDemux::default(); // Opening tag split across chunks. assert!(d.push("\nstep one"); - assert_eq!( - p1, - vec![ThinkPiece::Thinking("\nstep one".into())] - ); + assert_eq!(p1, vec![ThinkPiece::Thinking("\nstep one".into())]); let p2 = d.push(" continues\n\n## Answer\nHi"); - assert_eq!( - p3, - vec![ThinkPiece::Text("\n\n## Answer\nHi".into())] - ); + assert_eq!(p3, vec![ThinkPiece::Text("\n\n## Answer\nHi".into())]); assert!(d.finish().is_empty()); } #[test] fn think_tag_demux_plain_text_passthrough() { let mut d = ThinkTagDemux::default(); - assert_eq!( - d.push("hello "), - vec![ThinkPiece::Text("hello ".into())] - ); + assert_eq!(d.push("hello "), vec![ThinkPiece::Text("hello ".into())]); assert_eq!(d.push("world"), vec![ThinkPiece::Text("world".into())]); } @@ -4649,7 +4717,10 @@ mod tests { } sanitize_assistant_content(&mut content); assert_eq!(thinking, "\nThe user is asking what.\n"); - assert!(!thinking.contains("The userThe user"), "duplicated thinking: {thinking}"); + assert!( + !thinking.contains("The userThe user"), + "duplicated thinking: {thinking}" + ); assert_eq!(content, "Hello"); assert!(!content.contains("")); } @@ -4664,7 +4735,7 @@ mod tests { assert_eq!(s2, "Hi there"); } - #[test] + #[test] fn append_stream_fragment_handles_cumulative_and_delta() { // Pure incremental: repeated prefix tokens must NOT be dropped. let mut inc = String::new(); diff --git a/core/src/providers/streaming.rs b/core/src/providers/streaming.rs index 38fd338..a9915cf 100644 --- a/core/src/providers/streaming.rs +++ b/core/src/providers/streaming.rs @@ -53,6 +53,14 @@ fn token_count(value: &Value) -> Option { }) } +/// Hard cap on streamed tool-call `index` growth. A malicious or buggy gateway +/// can emit `index: 1e9` and OOM the process if the accumulator grows unbounded +/// (`while len <= idx { push }`). Parallel tool batches stay well under this. +pub(crate) const MAX_STREAM_TOOL_CALL_INDEX: usize = 64; + +/// Cap total argument bytes across all streamed tool calls in one turn. +pub(crate) const MAX_STREAM_TOOL_ARGS_BYTES: usize = 8 * 1024 * 1024; + /// Tool-call indices arrive as ints, floats (`1.0`), or quoted numbers on some /// gateways. Falling back to `0` silently merges parallel tool streams. fn tool_call_index(value: Option<&Value>) -> usize { @@ -81,6 +89,40 @@ fn tool_call_index(value: Option<&Value>) -> usize { .unwrap_or(0) as usize } +/// Ensure `tool_calls` has a slot at `idx` without unbounded allocation (H3). +pub(crate) fn ensure_stream_tool_slot( + tool_calls: &mut Vec, + idx: usize, +) -> Result<(), String> { + if idx >= MAX_STREAM_TOOL_CALL_INDEX { + return Err(format!( + "stream tool_call index {idx} exceeds cap ({MAX_STREAM_TOOL_CALL_INDEX})" + )); + } + while tool_calls.len() <= idx { + tool_calls.push(T::default()); + } + Ok(()) +} + +/// Append a tool-call argument fragment if total args stay under the byte cap. +pub(crate) fn append_stream_tool_args( + total_args_bytes: &mut usize, + acc_args: &mut String, + fragment: &str, +) -> Result<(), String> { + let add = fragment.len(); + let next = total_args_bytes.saturating_add(add); + if next > MAX_STREAM_TOOL_ARGS_BYTES { + return Err(format!( + "stream tool_call arguments exceed cap ({MAX_STREAM_TOOL_ARGS_BYTES} bytes)" + )); + } + *total_args_bytes = next; + acc_args.push_str(fragment); + Ok(()) +} + /// Some providers stream `function.arguments` as a JSON object/array instead of /// a string fragment. Convert those to a compact JSON string so the accumulator /// can still assemble a valid arguments payload. @@ -460,6 +502,34 @@ mod tests { } } + #[test] + fn ensure_stream_tool_slot_rejects_malicious_high_index() { + let mut slots: Vec = Vec::new(); + assert!(ensure_stream_tool_slot(&mut slots, 0).is_ok()); + assert_eq!(slots.len(), 1); + assert!(ensure_stream_tool_slot(&mut slots, MAX_STREAM_TOOL_CALL_INDEX - 1).is_ok()); + assert_eq!(slots.len(), MAX_STREAM_TOOL_CALL_INDEX); + let err = ensure_stream_tool_slot(&mut slots, MAX_STREAM_TOOL_CALL_INDEX).unwrap_err(); + assert!(err.contains("exceeds cap"), "{err}"); + // Must not allocate multi-GB for a huge index. + let before = slots.len(); + let _ = ensure_stream_tool_slot(&mut slots, usize::MAX / 2); + assert_eq!(slots.len(), before); + } + + #[test] + fn append_stream_tool_args_caps_total_bytes() { + let mut total = 0usize; + let mut acc = String::new(); + append_stream_tool_args(&mut total, &mut acc, "abc").unwrap(); + assert_eq!(total, 3); + assert_eq!(acc, "abc"); + let big = "x".repeat(MAX_STREAM_TOOL_ARGS_BYTES); + let err = append_stream_tool_args(&mut total, &mut acc, &big).unwrap_err(); + assert!(err.contains("exceed cap"), "{err}"); + assert_eq!(acc, "abc"); + } + #[test] fn function_call_finish_reason_maps_to_tool_calls() { let events = decode_openai_chunk(&json!({ diff --git a/core/src/subagent.rs b/core/src/subagent.rs index 560549e..034c6c9 100644 --- a/core/src/subagent.rs +++ b/core/src/subagent.rs @@ -1050,7 +1050,10 @@ pub fn execute( ) .await; if let Some(ref wt) = wt_path { - if outcome.ok { + // Cancelled runs report non-ok. Also refuse promote if the + // parent turn token fired (H13). run_single uses a child token + // for interrupt; abort paths already return Outcome::err. + if outcome.ok && !cancel.is_cancelled() { match crate::worktree::promote_worktree(&workspace, wt) { Ok(paths) if !paths.is_empty() => { emit( @@ -1606,7 +1609,7 @@ async fn run_agent_inner( loop { if cancel.is_cancelled() { - return Outcome::ok("[subagent aborted]"); + return Outcome::err("[subagent aborted]"); } // Poll intercom mailbox for orchestrator steer messages (peek + steer actions) // Drain orchestrator steer messages from the mailbox without @@ -1826,7 +1829,7 @@ async fn run_agent_inner( // destructive writes don't execute once cancelled. if cancel.is_cancelled() { emit_subagent_summary(sub_in, sub_out, sub_cached, &last_model); - return Outcome::ok("[subagent aborted]"); + return Outcome::err("[subagent aborted]"); } let id = call.id.clone(); let name = call.function.name.clone(); @@ -2078,7 +2081,7 @@ async fn dispatch_subagent_tool( return Outcome::err(format!("tool call '{}' was denied by the user", name)); } crate::ApprovalResult::Aborted => { - return Outcome::ok("[subagent aborted]"); + return Outcome::err("[subagent aborted]"); } } } @@ -3047,9 +3050,11 @@ async fn run_parallel( let all_ok = collected.iter().all(|(_, o)| o.ok); // Promote successful worktrees into the main workspace, then clean up. + // Never promote on cancel even if a child reported ok before the abort (H13). + let batch_cancelled = run_cancel.is_cancelled(); for (i, o) in &collected { if let Some(wt) = worktrees.get(*i).and_then(|w| w.as_ref()) { - if o.ok { + if o.ok && !batch_cancelled { match crate::worktree::promote_worktree(&workspace, wt) { Ok(paths) if !paths.is_empty() => { emit( @@ -3211,7 +3216,7 @@ async fn run_chain( for (step_i, step) in chain.iter().enumerate() { if run_cancel.is_cancelled() { - return Outcome::ok("[chain aborted]"); + return Outcome::err("[chain aborted]"); } // parallel group? if let Some(group) = step.get("parallel").and_then(|v| v.as_array()) { @@ -3987,4 +3992,28 @@ mod tests { assert!(crate::workspace::resolve(&workspace, "../secret.txt").is_err()); let _ = std::fs::remove_dir_all(&root); } + + #[test] + fn cancelled_subagent_outcome_is_not_ok() { + // H13: abort/cancel must return non-ok so worktree promote gates + // (`outcome.ok && !cancel`) never merge half-done trees. + let aborted = Outcome::err("[subagent aborted]"); + assert!(!aborted.ok); + assert_eq!(aborted.output, "[subagent aborted]"); + let chain = Outcome::err("[chain aborted]"); + assert!(!chain.ok); + } + + #[test] + fn promote_gate_requires_ok_and_not_cancelled() { + // Mirrors the single/parallel promote predicate used after run_single. + let ok = true; + let cancelled = true; + assert!( + !(ok && !cancelled), + "cancel must block promote even when ok" + ); + assert!(ok && !false, "successful non-cancelled run may promote"); + assert!(!(!true && !false), "failed run must not promote"); + } } diff --git a/core/src/tools.rs b/core/src/tools.rs index a98f214..fafc8ab 100644 --- a/core/src/tools.rs +++ b/core/src/tools.rs @@ -79,7 +79,16 @@ pub fn execute(name: &str, args: &Value, cfg: &Config) -> Outcome { _ => Outcome::err("edit requires a non-empty 'edits' array"), } } - "write_file" => write_file(s("path"), s("content"), cfg), + "write_file" => { + let path = s("path"); + // Require the `content` key (schema marks it required). Missing key + // used to silently write "" and truncate existing files (H5). + // Explicit empty string is allowed when the key is present. + match require_string_arg(args, "content") { + Ok(content) => write_file(path, content, cfg), + Err(e) => Outcome::err(e), + } + } "delete" => delete_path(s("path"), cfg), "rename" => rename_path(s("from"), s("to"), cfg), "mkdir" => mkdir_path(s("path"), cfg), @@ -125,6 +134,16 @@ impl Outcome { } } +/// Require a string argument key to be present. Null or non-string values fail. +/// Empty string is allowed when the key exists (explicit wipe). +fn require_string_arg<'a>(args: &'a Value, key: &str) -> Result<&'a str, String> { + match args.get(key) { + Some(Value::String(s)) => Ok(s.as_str()), + Some(_) => Err(format!("'{key}' must be a string")), + None => Err(format!("missing required '{key}'")), + } +} + // ---- file tools ---- /// Resolve a tool path against the workspace root. Approval::Never means @@ -2099,12 +2118,19 @@ fn bulk_write(args: &Value, cfg: &Config) -> Outcome { let mut ok = true; for (i, f) in files.iter().enumerate() { let path = f.get("path").and_then(|v| v.as_str()).unwrap_or(""); - let content = f.get("content").and_then(|v| v.as_str()).unwrap_or(""); if path.is_empty() { ok = false; lines.push(format!("[{i}] error: missing 'path'")); continue; } + let content = match require_string_arg(f, "content") { + Ok(c) => c, + Err(e) => { + ok = false; + lines.push(format!("[{i}] {path}: error: {e}")); + continue; + } + }; let r = write_file(path, content, cfg); if !r.ok { ok = false; @@ -4151,6 +4177,61 @@ mod tests { assert!(cfg.workspace.join(".git/config").exists()); } + #[test] + fn write_file_missing_content_errors_and_leaves_file() { + // H5: omit `content` → tool error; existing file must not be emptied. + let (_root, cfg) = tmp_ws(); + let path = cfg.workspace.join("keep.txt"); + fs::write(&path, "precious").unwrap(); + let o = execute("write_file", &json!({"path": "keep.txt"}), &cfg); + assert!(!o.ok, "missing content must fail: {}", o.output); + assert!( + o.output.contains("content"), + "error should mention content: {}", + o.output + ); + assert_eq!(fs::read_to_string(&path).unwrap(), "precious"); + } + + #[test] + fn write_file_explicit_empty_content_allowed() { + let (_root, cfg) = tmp_ws(); + let path = cfg.workspace.join("wipe.txt"); + fs::write(&path, "old").unwrap(); + let o = execute( + "write_file", + &json!({"path": "wipe.txt", "content": ""}), + &cfg, + ); + assert!(o.ok, "{}", o.output); + assert_eq!(fs::read_to_string(&path).unwrap(), ""); + } + + #[test] + fn bulk_write_missing_content_errors_and_skips_file() { + let (_root, cfg) = tmp_ws(); + let keep = cfg.workspace.join("keep.txt"); + fs::write(&keep, "safe").unwrap(); + let o = bulk_write( + &json!({"files":[ + {"path":"keep.txt"}, + {"path":"ok.txt","content":"hi"} + ]}), + &cfg, + ); + assert!(!o.ok, "{}", o.output); + assert!( + o.output.contains("content"), + "error should mention content: {}", + o.output + ); + assert_eq!(fs::read_to_string(&keep).unwrap(), "safe"); + assert_eq!( + fs::read_to_string(cfg.workspace.join("ok.txt")).unwrap(), + "hi" + ); + } + #[test] fn bulk_write_primitive_no_longer_blocks_restricted_paths() { // Mirrors write_file: bulk_write calls write_file, which no longer