From d9f7c829343e5c7acab0dc9f90183a00bc00d154 Mon Sep 17 00:00:00 2001 From: Xavier Yen Date: Thu, 27 Aug 2026 05:49:05 +0000 Subject: [PATCH] fix(agent): pair parallel tool calls by internal call id --- docs/ARCHITECTURE.md | 2 +- src/agent/runner.rs | 88 ++++-- src/event.rs | 9 + src/extras/acp/mod.rs | 43 ++- src/tests/mod.rs | 2 + src/tests/parallel_tool_call_tests.rs | 376 ++++++++++++++++++++++++++ src/ui/app.rs | 4 +- src/ui/event_handler.rs | 73 +++-- src/ui/mod.rs | 7 +- src/ui/state.rs | 64 ++++- 10 files changed, 611 insertions(+), 57 deletions(-) create mode 100644 src/tests/parallel_tool_call_tests.rs diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 079cd3ef..d0a101f0 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -162,7 +162,7 @@ reproducing code; focus on structure, relationships, and rationale. ## Data Flow - User input → InputEditor → event loop → agent.spawn_runner() -- Runner streams AgentEvent (Token, Reasoning, ToolCall, ToolResult, Done) +- Runner streams AgentEvent (Token, Reasoning, ToolCall, ToolResult, Done); ToolCall/ToolResult carry a correlation `call_id` (rig's `internal_call_id`) so a parallel batch pairs each result with its own call - Events rendered incrementally via Renderer::write_line() ## Design Decisions diff --git a/src/agent/runner.rs b/src/agent/runner.rs index c66485cf..1acd7e2c 100644 --- a/src/agent/runner.rs +++ b/src/agent/runner.rs @@ -1,3 +1,5 @@ +use std::collections::HashMap; + use compact_str::CompactString; use futures::StreamExt; use rig::agent::{Agent, MultiTurnStreamItem, StreamingResult}; @@ -338,7 +340,11 @@ where let retry_prompt = prompt.clone(); let retry_history: Vec = history.clone(); let mut tool_interactions: Vec = Vec::new(); - let mut last_tool_name: Option = None; + // In-flight calls by rig `internal_call_id`. A map, not a single + // slot: providers may stream a whole batch of parallel `ToolCall`s + // before any of their `ToolResult`s, so pairing by "most recent call" + // records the wrong name against a result. + let mut pending_tool_names: HashMap = HashMap::new(); let mut empty_response_count: u32 = 0; const MAX_EMPTY_RESPONSES: u32 = 3; // Overrides the next continuation message (bottom of the outer @@ -396,6 +402,8 @@ where }; loop { + // Entries orphaned by an abandoned turn must not outlive it. + pending_tool_names.clear(); while let Some(item) = stream.next().await { match item { Ok(MultiTurnStreamItem::StreamAssistantItem(content)) => { @@ -410,17 +418,22 @@ where .send(AgentEvent::Token(CompactString::from(text.text))) .await; } - StreamedAssistantContent::ToolCall { tool_call, .. } => { + StreamedAssistantContent::ToolCall { + tool_call, + internal_call_id, + } => { let tool_name = &tool_call.function.name; tracing::debug!( "agent tool start: name={}, args_len={}", tool_name, tool_call.function.arguments.to_string().len(), ); - last_tool_name = Some(tool_name.clone()); + pending_tool_names + .insert(internal_call_id.clone(), tool_name.clone()); tool_interactions.push(tool_call.clone().into()); let _ = event_tx .send(AgentEvent::ToolCall { + call_id: CompactString::from(internal_call_id), name: CompactString::from(tool_call.function.name), args: tool_call.function.arguments, }) @@ -431,10 +444,23 @@ where } Ok(MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult { tool_result, - .. + internal_call_id, })) => { - let tool_name = - CompactString::new(last_tool_name.take().unwrap_or_default()); + // Reachable on an abandoned turn; rig's `ToolResult` + // carries only call ids, so the name is unrecoverable + // and the consumer pairs on the id instead. + let tool_name = CompactString::new( + pending_tool_names + .remove(&internal_call_id) + .unwrap_or_else(|| { + tracing::warn!( + "tool result with no matching pending call \ + (internal_call_id={id})", + id = internal_call_id.escape_debug(), + ); + String::new() + }), + ); let mut output = String::new(); for c in tool_result.content.iter() { if let ToolResultContent::Text(t) = c { @@ -451,6 +477,7 @@ where ); let _ = event_tx .send(AgentEvent::ToolResult { + call_id: CompactString::from(internal_call_id), name: tool_name.clone(), output: CompactString::from(output), }) @@ -608,8 +635,9 @@ where #[cfg(feature = "hooks")] let mut tool_interactions: Vec = Vec::new(); let mut full_response = String::new(); - let mut last_tool_name: Option = None; - let mut last_tool_args: Option = None; + // In-flight calls by rig `internal_call_id`: (name, args). A map, not a + // pair of single slots — see the matching comment in `spawn_agent`. + let mut pending_calls: HashMap = HashMap::new(); // Unconditional (independent of `pure_stdout` and the `hooks` feature) // ordered record of this turn's completed tool call/result round trips, // returned to the caller (`dispatch_print`) for session persistence. See @@ -651,6 +679,8 @@ where while continue_turn { continue_turn = false; + // Entries orphaned by an abandoned turn must not outlive it. + pending_calls.clear(); loop { // Wait for the next stream item while staying available to the // subagent channel. `StreamExt::next` is cancel-safe (it only @@ -690,26 +720,41 @@ where let _ = std::io::Write::flush(&mut std::io::stderr()); } Ok(MultiTurnStreamItem::StreamAssistantItem( - StreamedAssistantContent::ToolCall { tool_call, .. }, + StreamedAssistantContent::ToolCall { + tool_call, + internal_call_id, + }, )) => { let name = tool_call.function.name.clone(); let args = tool_call.function.arguments.clone(); - last_tool_name = Some(name.clone()); - last_tool_args = Some(args.clone()); if pure_stdout { let summary = format_tool_args_summary(&args); println!("\n◈ {} {}", name, summary); let _ = std::io::Write::flush(&mut std::io::stdout()); } + pending_calls.insert(internal_call_id, (name, args)); #[cfg(feature = "hooks")] tool_interactions.push(tool_call.clone().into()); } Ok(MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult { tool_result, - .. + internal_call_id, })) => { - let name = last_tool_name.take().unwrap_or_default(); - let args = last_tool_args.take().unwrap_or(serde_json::Value::Null); + // Reachable: rig streams a result for a call that produced + // no `ToolCall` item when the turn is abandoned over an + // invalid tool call. The recorded pair stays + // self-consistent (a phantom call is recorded alongside + // it, see `startup.rs`); the name has to be empty because + // rig's `ToolResult` carries only call ids, not the name. + let (name, args) = + pending_calls.remove(&internal_call_id).unwrap_or_else(|| { + tracing::warn!( + "tool result with no matching pending call \ + (internal_call_id={id})", + id = internal_call_id.escape_debug(), + ); + (String::new(), serde_json::Value::Null) + }); let mut output = String::new(); for c in tool_result.content.iter() { if let ToolResultContent::Text(t) = c { @@ -731,12 +776,15 @@ where } let _ = std::io::Write::flush(&mut std::io::stdout()); } - // Anything still queued belongs to the call that just + // Attribute anything still queued to the call that just // finished: a subagent only runs inside its `task` call, // and every send completes before that call returns. The // `select!` above normally has them already, but a tool // that sends without ever yielding hands us its result in // the same poll, leaving them queued until here. + // Best-effort under a parallel batch: with several calls + // in flight the side channel carries no call id, so a + // sibling's result can claim the `task` call's subagents. #[cfg(feature = "subagents")] while let Ok(event) = subagent_rx.try_recv() { push_subagent_call(&mut pending_subagent_calls, event); @@ -824,8 +872,9 @@ where /// Vec` above, which carries the raw `rig` message types needed /// only for `Stop`-continuation replay. `dispatch_print` turns each of these /// into a `Session::add_tool_call` + `add_tool_result` pair (design.md -/// decision 5); relies on the single-threaded call/result pairing confirmed -/// in design.md's Open Questions (2.3). +/// decision 5); each round trip is paired by rig's `internal_call_id`, so a +/// parallel batch (every call streamed before the first result) records each +/// result against its own call. #[derive(Debug, Clone)] pub struct ToolInteraction { pub name: String, @@ -836,7 +885,10 @@ pub struct ToolInteraction { /// concurrency boundary: only this loop knows which main-agent call was /// in flight when each event arrived, and `dispatch_print` turns the /// nesting into `parent_call_id` once the enclosing call has an id. - /// Always empty for anything but a `task` call. + /// Empty for anything but a `task` call when that call runs alone; under a + /// parallel batch the attribution is best-effort, since queued subagent + /// calls attach to the batch's first-arriving result whichever call it + /// answers. #[cfg(feature = "subagents")] pub subagent_calls: Vec, } diff --git a/src/event.rs b/src/event.rs index 884cae87..1d737c06 100644 --- a/src/event.rs +++ b/src/event.rs @@ -5,10 +5,19 @@ pub enum AgentEvent { Token(CompactString), Reasoning(CompactString), ToolCall { + /// Rig's `internal_call_id` for this call: unique within the process + /// (unlike the provider-supplied `ToolCall.id`, which Gemini/Ollama + /// set to the function name). Providers may emit a whole batch of + /// parallel calls before any of their results, so consumers must pair + /// a `ToolResult` with its call by this id, never by "most recent + /// call". + call_id: CompactString, name: CompactString, args: serde_json::Value, }, ToolResult { + /// The [`AgentEvent::ToolCall::call_id`] this result answers. + call_id: CompactString, name: CompactString, output: CompactString, }, diff --git a/src/extras/acp/mod.rs b/src/extras/acp/mod.rs index 4d84466e..b2577dcc 100644 --- a/src/extras/acp/mod.rs +++ b/src/extras/acp/mod.rs @@ -8,6 +8,7 @@ use agent_client_protocol::schema::v1::*; use agent_client_protocol::{ Agent, ByteStreams, Client, ConnectTo, ConnectionTo, Dispatch, Responder, Role, Stdio, }; +use compact_str::CompactString; use tokio::sync::Mutex; use crate::cli::Cli; @@ -342,7 +343,11 @@ async fn run_prompt( .await; let mut rx = runner.event_rx; - let mut tool_call_id: Option = None; + // In-flight main-agent calls by `AgentEvent` id (rig's + // `internal_call_id`) to the ACP ToolCallId announced for them. A map, + // not a single slot: a parallel batch streams every `ToolCall` before + // the first `ToolResult`. + let mut tool_call_ids: HashMap = HashMap::new(); let mut final_response = String::new(); while let Some(event) = rx.recv().await { @@ -370,9 +375,13 @@ async fn run_prompt( tracing::warn!("ACP failed to send reasoning notification: {}", e); } } - AgentEvent::ToolCall { name, args } => { + AgentEvent::ToolCall { + call_id: event_id, + name, + args, + } => { let id = ToolCallId::new(uuid::Uuid::new_v4().to_string()); - tool_call_id = Some(id.clone()); + tool_call_ids.insert(event_id, id.clone()); let args_str = args.to_string(); let tool_call = ToolCall::new(id.clone(), name.to_string()) .raw_input(serde_json::from_str(&args_str).ok()); @@ -385,10 +394,16 @@ async fn run_prompt( } } AgentEvent::SubagentToolCall { name, args } => { + // Announce-only: subagent calls carry no correlating id, so + // they never receive a ToolCallUpdate. (Previously they + // hijacked the single pending slot, so the enclosing `task` + // call's result got attached to the subagent's entry.) + // Announced as already Completed, since nothing will ever + // update it out of the default Pending status. let id = ToolCallId::new(uuid::Uuid::new_v4().to_string()); - tool_call_id = Some(id.clone()); let args_str = args.to_string(); let tool_call = ToolCall::new(id.clone(), format!("[subagent] {}", name)) + .status(ToolCallStatus::Completed) .raw_input(serde_json::from_str(&args_str).ok()); let notif = SessionNotification::new( session_id.clone(), @@ -398,10 +413,22 @@ async fn run_prompt( tracing::warn!("ACP failed to send subagent tool call notification: {}", e); } } - AgentEvent::ToolResult { output, .. } => { - let id = tool_call_id - .take() - .unwrap_or_else(|| ToolCallId::new(uuid::Uuid::new_v4().to_string())); + AgentEvent::ToolResult { + call_id: event_id, + output, + .. + } => { + // No announced ToolCall to update: an update carrying a + // ToolCallId the client was never told about is worse than + // silence, so drop it. + let Some(id) = tool_call_ids.remove(&event_id) else { + tracing::warn!( + "ACP tool result with no announced tool call (id={}); \ + skipping update", + event_id.escape_debug(), + ); + continue; + }; let fields = ToolCallUpdateFields::new() .status(ToolCallStatus::Completed) .content(vec![ToolCallContent::from(ContentBlock::Text( diff --git a/src/tests/mod.rs b/src/tests/mod.rs index a666ddf4..1035df39 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -67,6 +67,8 @@ mod multimodal_tests; #[cfg(test)] mod normalize_tests; #[cfg(test)] +mod parallel_tool_call_tests; +#[cfg(test)] mod paste_burst_tests; #[cfg(test)] mod picker_tests; diff --git a/src/tests/parallel_tool_call_tests.rs b/src/tests/parallel_tool_call_tests.rs new file mode 100644 index 00000000..cd69d269 --- /dev/null +++ b/src/tests/parallel_tool_call_tests.rs @@ -0,0 +1,376 @@ +//! Regression proof that a parallel tool-call batch (all `ToolCall`s +//! streamed before any of their `ToolResult`s, which is how rig 0.40 drives +//! providers' parallel tool use) is recorded with each result paired to its +//! own call. Before pairing was keyed by rig's `internal_call_id`, both +//! `run_print` and `spawn_agent` tracked the batch in a single +//! most-recent-call slot, so with N>1 calls in flight the first result was +//! recorded under the *last* call's name/args and later results under an +//! empty name — corrupting the session-JSON evidence channel. +//! +//! Uses the same fake-model + real-registered-tool setup as +//! `headless_tool_record_tests.rs` (see that file's header), with two +//! distinct tools so a mispairing changes observable names, not just args. +//! +//! The TUI keeps its own copy of that pairing in `AgentRunState`, driven by +//! the same events; the second half of this file covers that bookkeeping and +//! the unmatched-result fallback directly, since neither is reachable from the +//! fake model (rig only streams an unmatched result when a hook stack skips an +//! invalid tool call, and zerostack installs no rig hooks). + +use rig::agent::AgentBuilder; +use rig::tool::Tool; +use serde::Deserialize; + +use crate::agent::runner::run_print; +use crate::agent::tools::ToolError; +use crate::event::AgentEvent; +use crate::retry::RetryConfig; +use crate::session::{Session, ToolRecord}; +use crate::tests::fake_model::{FakeModel, MockCompletionModel, MockStreamEvent}; +use crate::ui::event_handler::resolve_tool_result_call_id; +use crate::ui::state::AgentRunState; + +#[derive(Debug, Deserialize)] +struct TextArgs { + text: String, +} + +/// A tool event as `spawn_agent` emitted it, kept in arrival order so a test +/// can assert on the batch shape as well as on the pairing. +#[derive(Debug)] +enum ToolEvent { + Call { + id: compact_str::CompactString, + name: String, + args: serde_json::Value, + }, + Result { + id: compact_str::CompactString, + name: String, + output: String, + }, +} + +struct EchoTool; + +impl Tool for EchoTool { + const NAME: &'static str = "echo"; + + type Error = ToolError; + type Args = TextArgs; + type Output = String; + + fn description(&self) -> String { + "Echoes the given text back.".to_string() + } + + fn parameters(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { "text": { "type": "string" } }, + "required": ["text"] + }) + } + + async fn call(&self, args: TextArgs) -> Result { + Ok(format!("echoed: {}", args.text)) + } +} + +struct ReverseTool; + +impl Tool for ReverseTool { + const NAME: &'static str = "reverse"; + + type Error = ToolError; + type Args = TextArgs; + type Output = String; + + fn description(&self) -> String { + "Reverses the given text.".to_string() + } + + fn parameters(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { "text": { "type": "string" } }, + "required": ["text"] + }) + } + + async fn call(&self, args: TextArgs) -> Result { + Ok(args.text.chars().rev().collect()) + } +} + +/// One model turn issuing two tool calls in a single batch, then a plain-text +/// closing turn. +fn two_call_model() -> FakeModel { + MockCompletionModel::from_stream_turns(vec![ + vec![ + MockStreamEvent::tool_call("call-1", "echo", serde_json::json!({ "text": "alpha" })), + MockStreamEvent::tool_call("call-2", "reverse", serde_json::json!({ "text": "beta" })), + MockStreamEvent::final_response_with_default_usage(), + ], + vec![ + MockStreamEvent::text("done".to_string()), + MockStreamEvent::final_response_with_default_usage(), + ], + ]) +} + +#[tokio::test] +async fn run_print_pairs_each_result_with_its_own_call_in_a_parallel_batch() { + // `run_print` reaches process-global wiring (the hooks Stop dispatcher, + // the subagent event sender); serialize against every other `run_print` + // test so none of them clobber each other's. + let _run_print_guard = crate::tests::fake_model::run_print_guard::acquire(); + + let model = two_call_model(); + let agent = AgentBuilder::new(model) + .tool(EchoTool) + .tool(ReverseTool) + .default_max_turns(2) + .build(); + + let outcome = run_print( + &agent, + "please echo and reverse", + false, + &RetryConfig::default(), + Vec::new(), + #[cfg(feature = "hooks")] + None, + ) + .await + .expect("run_print should succeed against the fake model"); + + assert_eq!(outcome.tool_interactions.len(), 2); + + let echo = outcome + .tool_interactions + .iter() + .find(|i| i.name == "echo") + .expect( + "an interaction named 'echo' must be recorded (an empty or \ + mispaired name means single-slot tracking regressed)", + ); + assert_eq!(echo.args, serde_json::json!({ "text": "alpha" })); + assert_eq!(echo.output, "echoed: alpha"); + + let reverse = outcome + .tool_interactions + .iter() + .find(|i| i.name == "reverse") + .expect("an interaction named 'reverse' must be recorded"); + assert_eq!(reverse.args, serde_json::json!({ "text": "beta" })); + assert_eq!(reverse.output, "ateb"); +} + +#[tokio::test] +async fn spawn_agent_events_pair_each_result_with_its_own_call_by_id() { + // `spawn_agent` sets the same process-global subagent event sender as + // `run_print`; hold the shared guard for the same reason. + let _run_print_guard = crate::tests::fake_model::run_print_guard::acquire(); + + let model = two_call_model(); + let agent = AgentBuilder::new(model) + .tool(EchoTool) + .tool(ReverseTool) + .default_max_turns(2) + .build(); + + let mut runner = crate::agent::runner::spawn_agent( + agent, + "please echo and reverse".to_string(), + Vec::new(), + RetryConfig::default(), + #[cfg(feature = "hooks")] + None, + ); + + // One ordered log rather than a call vector and a result vector: the + // arrival order is itself under test (see the batch-shape assertion + // below), and two vectors discard it. + let mut log: Vec = Vec::new(); + while let Some(event) = runner.event_rx.recv().await { + match event { + AgentEvent::ToolCall { + call_id, + name, + args, + } => log.push(ToolEvent::Call { + id: call_id, + name: name.to_string(), + args, + }), + AgentEvent::ToolResult { + call_id, + name, + output, + } => log.push(ToolEvent::Result { + id: call_id, + name: name.to_string(), + output: output.to_string(), + }), + AgentEvent::Done { .. } | AgentEvent::Error(_) => break, + _ => {} + } + } + + let first_result = log + .iter() + .position(|e| matches!(e, ToolEvent::Result { .. })) + .expect("the fake model's two calls must produce results"); + let calls_before_first_result = log[..first_result] + .iter() + .filter(|e| matches!(e, ToolEvent::Call { .. })) + .count(); + assert_eq!( + calls_before_first_result, 2, + "this test's premise is the parallel-batch shape: both ToolCall events \ + must arrive before the first ToolResult. If rig switched to strict \ + call/result interleaving, the pairing assertions below would pass \ + without guarding anything" + ); + + let calls: Vec<(compact_str::CompactString, String, serde_json::Value)> = log + .iter() + .filter_map(|e| match e { + ToolEvent::Call { id, name, args } => Some((id.clone(), name.clone(), args.clone())), + ToolEvent::Result { .. } => None, + }) + .collect(); + let results: Vec<(compact_str::CompactString, String, String)> = log + .iter() + .filter_map(|e| match e { + ToolEvent::Result { id, name, output } => { + Some((id.clone(), name.clone(), output.clone())) + } + ToolEvent::Call { .. } => None, + }) + .collect(); + + assert_eq!(calls.len(), 2); + assert_eq!(results.len(), 2); + assert_ne!(calls[0].0, calls[1].0, "event ids must be unique per call"); + + for (result_id, result_name, output) in &results { + let (_, call_name, args) = calls + .iter() + .find(|(call_id, _, _)| call_id == result_id) + .expect("every ToolResult event must reference a ToolCall event's id"); + assert_eq!( + result_name, call_name, + "a ToolResult must carry the name of the call it answers" + ); + let expected = match call_name.as_str() { + "echo" => format!("echoed: {}", args["text"].as_str().unwrap()), + "reverse" => args["text"].as_str().unwrap().chars().rev().collect(), + other => panic!("unexpected tool name {other}"), + }; + assert_eq!(output, &expected); + } +} + +#[test] +fn agent_run_state_pairs_each_result_of_a_two_call_batch() { + let mut run = AgentRunState::default(); + run.push_pending_tool_call("call-1".into(), 7); + run.push_pending_tool_call("call-2".into(), 8); + + // Results may come back in either order; each must find its own call. + assert_eq!(run.take_pending_tool_call("call-2"), Some(8)); + assert_eq!(run.take_pending_tool_call("call-1"), Some(7)); + assert_eq!(run.take_pending_tool_call("call-1"), None); + assert!(run.pending_tool_calls.is_empty()); +} + +#[cfg(any(feature = "subagents", feature = "acp"))] +#[test] +fn newest_pending_tool_call_falls_back_to_the_next_newest() { + let mut run = AgentRunState::default(); + assert_eq!(run.newest_pending_tool_call(), None); + + run.push_pending_tool_call("call-1".into(), 7); + run.push_pending_tool_call("call-2".into(), 8); + assert_eq!(run.newest_pending_tool_call(), Some(8)); + + // Once the newest call's result lands it is no longer a candidate parent + // for a subagent call, but its still-running sibling is. + run.take_pending_tool_call("call-2"); + assert_eq!(run.newest_pending_tool_call(), Some(7)); + + run.take_pending_tool_call("call-1"); + assert_eq!(run.newest_pending_tool_call(), None); +} + +#[test] +fn pushing_a_duplicate_id_replaces_the_stale_entry() { + let mut run = AgentRunState::default(); + run.push_pending_tool_call("call-1".into(), 7); + run.push_pending_tool_call("call-2".into(), 8); + run.push_pending_tool_call("call-1".into(), 9); + + // One entry per id, and the live call wins — including for recency. + assert_eq!(run.pending_tool_calls.len(), 2); + #[cfg(any(feature = "subagents", feature = "acp"))] + assert_eq!(run.newest_pending_tool_call(), Some(9)); + assert_eq!(run.take_pending_tool_call("call-1"), Some(9)); + assert_eq!(run.take_pending_tool_call("call-2"), Some(8)); +} + +#[test] +fn clearing_drops_calls_stranded_by_a_teardown() { + let mut run = AgentRunState::default(); + run.push_pending_tool_call("call-1".into(), 7); + run.push_pending_tool_call("call-2".into(), 8); + + run.clear_pending_tool_calls(); + + assert!(run.pending_tool_calls.is_empty()); + assert_eq!(run.take_pending_tool_call("call-1"), None); + // The decisive one: a stranded entry would otherwise parent the respawned + // turn's first subagent call to the aborted turn's call. + #[cfg(any(feature = "subagents", feature = "acp"))] + assert_eq!(run.newest_pending_tool_call(), None); +} + +#[test] +fn an_unmatched_tool_result_gets_its_own_call_instead_of_call_zero() { + let mut session = Session::new("anthropic", "claude-test", 200_000, ""); + let mut run = AgentRunState::default(); + + // Session call ids start at 0, so the first real call owns id 0 — the + // value an unmatched result must never be linked to. + let real_call = session.add_tool_call("read", &serde_json::json!({ "path": "a.txt" })); + assert_eq!(real_call, 0); + run.push_pending_tool_call("call-1".into(), real_call); + + let orphan = resolve_tool_result_call_id(&mut run, &mut session, "call-ghost", "grep"); + assert_ne!(orphan, real_call); + session.add_tool_result(orphan, "grep", "no matches"); + + // The real call's record is untouched... + let calls: Vec<&ToolRecord> = session + .messages + .iter() + .filter_map(|m| m.tool.as_ref()) + .filter(|r| matches!(r, ToolRecord::Call { .. })) + .collect(); + assert!(matches!( + calls.first(), + Some(ToolRecord::Call { id, name, args }) + if *id == real_call && name == "read" && args["path"] == "a.txt" + )); + // ...and the orphan result answers a synthesized call of its own, with no + // args because none were ever streamed. + assert!(matches!( + calls.get(1), + Some(ToolRecord::Call { id, name, args }) + if *id == orphan && name == "grep" && args.is_null() + )); + + // The pending real call is still in flight and still findable. + assert_eq!(run.take_pending_tool_call("call-1"), Some(real_call)); +} diff --git a/src/ui/app.rs b/src/ui/app.rs index e791a655..9b46e266 100644 --- a/src/ui/app.rs +++ b/src/ui/app.rs @@ -956,7 +956,7 @@ impl<'a> App<'a> { async fn handle_agent_event(&mut self, event: AgentEvent) -> anyhow::Result<()> { match &event { - AgentEvent::ToolCall { name, args } => { + AgentEvent::ToolCall { name, args, .. } => { if self.run.turn_trace.len() < TURN_TRACE_MAX { self.run .turn_trace @@ -1126,6 +1126,7 @@ impl<'a> App<'a> { self.run.agent_rx = None; self.run.turn_trace.clear(); self.run.awaiting_compaction_relief = false; + self.run.clear_pending_tool_calls(); self.run.pending_inputs.clear(); #[cfg(feature = "loop")] if let Some(ref mut ls) = self.chain.loop_state { @@ -2225,6 +2226,7 @@ impl<'a> App<'a> { ss.send_stop(); } self.run.agent_rx = None; + self.run.clear_pending_tool_calls(); } } None diff --git a/src/ui/event_handler.rs b/src/ui/event_handler.rs index 204845d3..c882c5da 100644 --- a/src/ui/event_handler.rs +++ b/src/ui/event_handler.rs @@ -98,7 +98,11 @@ pub async fn handle_agent_event( renderer.render_viewport()?; run.agent_line_started = true; } - AgentEvent::ToolCall { name, args } => { + AgentEvent::ToolCall { + call_id: event_id, + name, + args, + } => { run.was_reasoning = false; if run.agent_line_started { renderer.write_line("", Color::White)?; @@ -106,7 +110,8 @@ pub async fn handle_agent_event( } run.response_buf.clear(); run.response_start_block = None; - run.pending_tool_call_id = Some(ui.session.add_tool_call(&name, &args)); + let call_id = ui.session.add_tool_call(&name, &args); + run.push_pending_tool_call(event_id, call_id); save_session_if_enabled(ui.session, ui.cli, renderer)?; let line = format!( "◈ {}", @@ -116,15 +121,15 @@ pub async fn handle_agent_event( } #[cfg(any(feature = "subagents", feature = "acp"))] AgentEvent::SubagentToolCall { name, args } => { - // Peeked, not taken: the enclosing `task` call's `ToolResult` is - // still to come and owns the consuming `take()`. Subagent events - // are sent from the task the `task` tool spawned, but they reach - // this handler through the same mpsc channel as the main agent's - // own events, and the subagent finishes before the `task` tool - // returns — so they always arrive between that call's `ToolCall` - // and `ToolResult`, with its id still pending here. + // Subagent events are sent from the task the `task` tool spawned, + // but they reach this handler through the same mpsc channel as + // the main agent's own events, and the subagent finishes before + // the `task` tool returns — so they arrive between that call's + // `ToolCall` and `ToolResult`, while its id is still pending here. + // (Best-effort under a parallel batch; see + // `AgentRunState::newest_pending_tool_call`.) ui.session - .add_subagent_tool_call(run.pending_tool_call_id, &name, &args); + .add_subagent_tool_call(run.newest_pending_tool_call(), &name, &args); save_session_if_enabled(ui.session, ui.cli, renderer)?; let line = format!( "⌥ {}", @@ -132,15 +137,12 @@ pub async fn handle_agent_event( ); renderer.write_line(&sanitize_output(&line), C_TOOL)?; } - AgentEvent::ToolResult { name, output } => { - let call_id = run.pending_tool_call_id.take().unwrap_or_else(|| { - tracing::warn!( - "ToolResult for {name} arrived with no pending ToolCall id; \ - linking to 0 (the agent event stream is expected to be strictly \ - sequential, so this should not happen)" - ); - 0 - }); + AgentEvent::ToolResult { + call_id: event_id, + name, + output, + } => { + let call_id = resolve_tool_result_call_id(run, ui.session, &event_id, &name); ui.session.add_tool_result(call_id, &name, &output); save_session_if_enabled(ui.session, ui.cli, renderer)?; if name == "todo_write" { @@ -301,12 +303,41 @@ pub async fn handle_agent_event( run.agent_line_started = false; run.response_buf.clear(); run.response_start_block = None; + // A mid-stream error strands whatever was in flight. + run.clear_pending_tool_calls(); save_session_if_enabled(ui.session, ui.cli, renderer)?; } } Ok(()) } +/// Session call id a `ToolResult` event's output belongs to. +/// +/// Normally the matching `ToolCall` event left a pending entry. rig also has a +/// path that streams a result for a call that never produced a `ToolCall` +/// stream item (an invalid tool call the hook stack skips, abandoning the +/// turn), so an unmatched result is possible. Session call id 0 is a real call +/// (`Session::next_tool_call_id` starts there), so such a result must not be +/// linked to it — that would rewrite the first call's recorded evidence. +/// Instead a call is synthesized for it, with no arguments because none were +/// ever streamed, keeping the recorded pair self-consistent. +pub(crate) fn resolve_tool_result_call_id( + run: &mut AgentRunState, + session: &mut Session, + id: &str, + name: &str, +) -> u64 { + if let Some(call_id) = run.take_pending_tool_call(id) { + return call_id; + } + tracing::warn!( + "ToolResult for {name} arrived with no matching pending ToolCall (id={escaped}); \ + recording it against a synthesized orphan call", + escaped = id.escape_debug(), + ); + session.add_tool_call(name, &serde_json::Value::Null) +} + fn save_session_if_enabled( session: &Session, cli: &Cli, @@ -429,6 +460,10 @@ async fn handle_agent_done( ss.send_stop(); } run.agent_rx = None; + // A clean turn has consumed every entry already; this covers a turn that + // ended with calls still open, and must precede the /loop respawn below so + // the next iteration starts empty. + run.clear_pending_tool_calls(); #[cfg(feature = "loop")] if let Some(ls) = chain.loop_state.as_mut() diff --git a/src/ui/mod.rs b/src/ui/mod.rs index 5ad47c5e..505f41d6 100644 --- a/src/ui/mod.rs +++ b/src/ui/mod.rs @@ -1,5 +1,5 @@ pub(crate) mod app; -mod event_handler; +pub(crate) mod event_handler; pub(crate) mod events; pub(crate) mod feed; pub(crate) mod input; @@ -627,6 +627,10 @@ pub(crate) async fn mid_turn_compact_and_respawn( run.is_running = false; run.agent_rx = None; run.was_reasoning = false; + // The aborted run's in-flight calls never get their results; the respawn + // below reuses this state, so a stranded entry would be paired with the + // new turn's first call. + run.clear_pending_tool_calls(); // 2. Record progress so far. `turn_trace` is a capped/truncated digest, so // this is best-effort continuity, paired with any partial response text. @@ -727,6 +731,7 @@ pub(crate) fn stop_turn_context_exhausted( run.turn_trace.clear(); run.response_buf.clear(); run.response_start_block = None; + run.clear_pending_tool_calls(); if let Some(ss) = ui.status_signals.as_ref() { ss.send_stop(); } diff --git a/src/ui/state.rs b/src/ui/state.rs index f2ad5cb0..de562674 100644 --- a/src/ui/state.rs +++ b/src/ui/state.rs @@ -4,6 +4,7 @@ use std::collections::VecDeque; +use compact_str::CompactString; use tokio::sync::mpsc; use crate::cli::Cli; @@ -138,16 +139,61 @@ pub(crate) struct AgentRunState { pub response_start_block: Option, pub pending_send: Option, pub was_reasoning: bool, - pub turn_trace: Vec, + pub turn_trace: Vec, pub awaiting_compaction_relief: bool, - /// The id `Session::add_tool_call` stamped for the most recent `ToolCall` - /// event, held until the matching `ToolResult` event consumes it via - /// `Session::add_tool_result`. The agent event stream is strictly - /// sequential (one tool call's result always arrives before the next - /// call starts — see design.md's single-threaded-execution finding), so a - /// single pending slot is sufficient; it never needs to hold more than - /// one id at a time. - pub pending_tool_call_id: Option, + /// In-flight `ToolCall` events in arrival order: each event's id (rig's + /// `internal_call_id`) paired with the id `Session::add_tool_call` stamped + /// for it, removed when the matching `ToolResult` event arrives. Several + /// are pending at once because providers may emit a whole batch of + /// parallel tool calls (every `ToolCall` event first, then their results), + /// so this cannot be a single slot; arrival order is kept so + /// [`newest_pending_tool_call`](AgentRunState::newest_pending_tool_call) + /// can answer. Reach it only through the methods below, which maintain + /// both invariants (unique ids, newest last). + pub pending_tool_calls: Vec<(CompactString, u64)>, +} + +impl AgentRunState { + /// Record `session_call_id` as the pending call for event id `id`. A + /// repeated id violates rig's uniqueness guarantee; the newer call is the + /// live one, so it replaces the stale entry and becomes the newest. + pub(crate) fn push_pending_tool_call(&mut self, id: CompactString, session_call_id: u64) { + if let Some(pos) = self.pending_tool_calls.iter().position(|(k, _)| *k == id) { + tracing::warn!( + "ToolCall id={escaped} was already in flight; replacing its pending entry \ + (rig internal_call_ids are expected to be unique)", + escaped = id.escape_debug(), + ); + self.pending_tool_calls.remove(pos); + } + self.pending_tool_calls.push((id, session_call_id)); + } + + /// Consume the session call id recorded for event id `id`, if any. + pub(crate) fn take_pending_tool_call(&mut self, id: &str) -> Option { + let pos = self.pending_tool_calls.iter().position(|(k, _)| k == id)?; + Some(self.pending_tool_calls.remove(pos).1) + } + + /// Session call id of the newest call still awaiting its result, or `None` + /// when nothing is in flight. `SubagentToolCall` events attach to it: the + /// subagent side channel carries no call id, so this is a best-effort + /// heuristic — under a parallel batch containing a `task` call it can + /// attribute a subagent call to a sibling. (`run_print` records the + /// opposite heuristic, attaching queued subagent calls to the batch's + /// first-arriving result; both stay best-effort until the side channel + /// carries a call id.) + #[cfg(any(feature = "subagents", feature = "acp"))] + pub(crate) fn newest_pending_tool_call(&self) -> Option { + self.pending_tool_calls.last().map(|&(_, call_id)| call_id) + } + + /// Drop every in-flight entry. Called from run teardown: a clean turn + /// consumes them all, but an abort or a mid-stream error strands some, and + /// a stranded entry must not be visible to the next run. + pub(crate) fn clear_pending_tool_calls(&mut self) { + self.pending_tool_calls.clear(); + } } /// What happens when the current run finishes: chained prompts, dot-prompt