From d3658b29a4f5d31ece1e6b18cd7d5bbc755eda49 Mon Sep 17 00:00:00 2001 From: "claude (aiteam)" Date: Tue, 1 Sep 2026 17:42:56 -0400 Subject: [PATCH] openai-chat streaming: emit ToolCallDelta so the stall watchdog sees tool-call progress (mu-b82rr) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The openai-chat provider accumulated tool-call fragments silently and surfaced them only in Done — a v1 shortcut from before the loop's stall watchdog (mu-197pd) learned to credit ToolCallDelta bytes. A tool call whose arguments stream longer than STREAM_STALL_SECS with no text or reasoning deltas (a large file written via one write call on a local lane) was misread as a dead connection and killed at exactly 300s. Emit a ToolCallDelta per fragment chunk (one-event-per-chunk slot preserved; Done-side assembly unchanged) and update the streaming tests to the new event counts, plus a regression test shaped like the failing capture. --- crates/mu-ai/src/providers/openrouter.rs | 30 +++-- .../mu-ai/src/providers/openrouter_tests.rs | 103 +++++++++++++++++- 2 files changed, 120 insertions(+), 13 deletions(-) diff --git a/crates/mu-ai/src/providers/openrouter.rs b/crates/mu-ai/src/providers/openrouter.rs index 8c7828d7..c4bd03e9 100644 --- a/crates/mu-ai/src/providers/openrouter.rs +++ b/crates/mu-ai/src/providers/openrouter.rs @@ -1010,7 +1010,12 @@ async fn next_event(mut state: StreamState) -> Option<(ProviderEvent, StreamStat } } } - // Tool call delta(s)? + // Tool call delta(s)? Accumulate for Done-side assembly AND + // stream a ToolCallDelta so the loop's stall watchdog counts + // the bytes (mu-b82rr): a large file written via one tool call + // streams arguments for minutes with no text/reasoning deltas + // at all, and an unemitted fragment is invisible progress the + // watchdog misreads as a dead connection. if let Some(deltas) = choice.delta.tool_calls { for tc_delta in deltas { let entry = state.tool_calls.entry(tc_delta.index).or_insert_with(|| { @@ -1021,17 +1026,29 @@ async fn next_event(mut state: StreamState) -> Option<(ProviderEvent, StreamStat if let Some(id) = tc_delta.id { entry.id = id; } + let mut name_delta = None; + let mut arguments_delta = None; if let Some(func) = tc_delta.function { if let Some(name) = func.name { + name_delta = Some(name.clone()); entry.name = name; } if let Some(args) = func.arguments { + arguments_delta = Some(args.clone()); entry.args_json.push_str(&args); } } - // v1: don't emit ProviderEvent::ToolCallDelta; - // the loop ignores it. Final tool calls are - // surfaced in Done. + if emitted_event.is_none() + && (name_delta.is_some() || arguments_delta.is_some()) + { + // Continuation fragments may carry no id; the loop + // uses it only for status display. + emitted_event = Some(ProviderEvent::ToolCallDelta { + id: entry.id.clone(), + name_delta, + arguments_delta, + }); + } } } // finish_reason landed? @@ -1043,9 +1060,8 @@ async fn next_event(mut state: StreamState) -> Option<(ProviderEvent, StreamStat if let Some(event) = emitted_event { return Some((event, state)); } - // No emittable event for this chunk (e.g. it was a delta - // with only tool_calls or a finish_reason). Loop and pull - // the next SSE event. + // No emittable event for this chunk (e.g. it carried only a + // finish_reason or usage). Loop and pull the next SSE event. } } diff --git a/crates/mu-ai/src/providers/openrouter_tests.rs b/crates/mu-ai/src/providers/openrouter_tests.rs index cb9eace7..d7543f87 100644 --- a/crates/mu-ai/src/providers/openrouter_tests.rs +++ b/crates/mu-ai/src/providers/openrouter_tests.rs @@ -1034,9 +1034,16 @@ async fn b8_sse_tool_call_accumulation() { v }; - // Just one Done event (we don't emit ToolCallDelta during streaming in v1). - assert_eq!(events.len(), 1); - let done = match events.into_iter().next().unwrap() { + // One ToolCallDelta per fragment chunk (mu-b82rr: the stall watchdog + // counts these bytes), then Done. + assert_eq!(events.len(), 4, "got {events:?}"); + for e in &events[..3] { + assert!( + matches!(e, ProviderEvent::ToolCallDelta { .. }), + "expected ToolCallDelta, got {e:?}" + ); + } + let done = match events.into_iter().nth(3).unwrap() { ProviderEvent::Done(msg) => msg, other => panic!("expected Done, got {other:?}"), }; @@ -1075,9 +1082,14 @@ async fn b9_sse_mixed_text_and_tool_call() { events.push(e); } - // 1 TextDelta + 1 Done. - assert_eq!(events.len(), 2); - let done = match events.into_iter().nth(1).unwrap() { + // 1 TextDelta + 1 ToolCallDelta (mu-b82rr) + 1 Done. + assert_eq!(events.len(), 3, "got {events:?}"); + assert!( + matches!(&events[1], ProviderEvent::ToolCallDelta { .. }), + "expected ToolCallDelta, got {:?}", + events[1] + ); + let done = match events.into_iter().nth(2).unwrap() { ProviderEvent::Done(msg) => msg, other => panic!("expected Done, got {other:?}"), }; @@ -1288,3 +1300,82 @@ mod live_tests { ); } } + +#[tokio::test] +async fn sse_tool_call_fragments_emit_toolcalldelta_mu_b82rr() { + // mu-b82rr: a tool call whose arguments stream for minutes (large file + // write on a slow lane) used to emit NO events between the last text + // delta and Done — the loop's stall watchdog counted zero bytes and + // killed the live connection at STREAM_STALL_SECS. Every argument + // fragment must surface as a ToolCallDelta; Done still assembles the + // complete call. + let raw = concat!( + r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_w1","function":{"name":"write","arguments":"{\"path\":\"a.h"}}]}}]}"#, + "\n\n", + r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"tml\",\"content\":\"\"}"}}]}}]}"#, + "\n\n", + r#"data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}"#, + "\n\n", + r#"data: [DONE]"#, + "\n\n", + ); + let bytes = futures::stream::iter(vec![Ok::<_, std::io::Error>(Bytes::copy_from_slice( + raw.as_bytes(), + ))]); + let (_tx, rx) = tokio::sync::oneshot::channel(); + let mut stream = test_events_stream(bytes, rx); + + let mut events = Vec::new(); + while let Some(e) = stream.next().await { + events.push(e); + } + + // 2 ToolCallDelta (one per fragment chunk) + 1 Done. + assert_eq!(events.len(), 3, "got {events:?}"); + match &events[0] { + ProviderEvent::ToolCallDelta { + id, + name_delta, + arguments_delta, + } => { + assert_eq!(id, "call_w1"); + assert_eq!(name_delta.as_deref(), Some("write")); + assert_eq!(arguments_delta.as_deref(), Some("{\"path\":\"a.h")); + } + other => panic!("expected ToolCallDelta, got {other:?}"), + } + match &events[1] { + ProviderEvent::ToolCallDelta { + id, + name_delta, + arguments_delta, + } => { + // Continuation fragment: id already known from the builder. + assert_eq!(id, "call_w1"); + assert!(name_delta.is_none()); + assert_eq!( + arguments_delta.as_deref(), + Some("tml\",\"content\":\"\"}") + ); + } + other => panic!("expected ToolCallDelta, got {other:?}"), + } + match &events[2] { + ProviderEvent::Done(msg) => { + assert_eq!(msg.stop_reason, StopReason::ToolUse); + assert_eq!(msg.content.len(), 1, "got {:?}", msg.content); + match &msg.content[0] { + ContentBlock::ToolCall(tc) => { + assert_eq!(tc.id, "call_w1"); + assert_eq!(tc.name, "write"); + assert_eq!( + tc.arguments.as_value().get("path").and_then(|v| v.as_str()), + Some("a.html") + ); + } + other => panic!("expected ToolCall, got {other:?}"), + } + } + other => panic!("expected Done, got {other:?}"), + } +}