From abc79e772d1c69d8a0cfc710d76baf5d539187fc Mon Sep 17 00:00:00 2001 From: Atish Patel Date: Wed, 29 Jul 2026 10:39:11 -0500 Subject: [PATCH 1/2] feat(agent): make Gemini and other MLflow-route models usable through databricks_v2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Non-benchmark changes lifted off benchmark/harness-accounting-and-solo so they can land independently. Two defects made every non-Claude, non-GPT-5 model on the Databricks MLflow route (databricks_v2) unusable in an agent loop; both live only in openai_body/parse_openai, which the luna (Responses) and opus (Anthropic) conditions never reach, so this is inert for models already in use. D1 — dropped thought signatures (fatal). Gemini returns a `thoughtSignature` on every tool call and requires it echoed back verbatim, as a sibling of `function` (nesting it inside `function{}` 400s with the same error as omitting it). buzz reserialized calls as `{id, type, function}` only, dropping the field, so the first tool call of every run failed with "Function call is missing a thought_signature". ToolCall now carries `provider_extra: Map`; parse_openai captures every unmodelled top-level key and openai_body re-emits them beside `function`, preserving position. Keeping whatever we did not model (rather than naming `thoughtSignature`) means the next provider with an opaque per-call token needs no change. The Responses and Anthropic replay shapes are fully modelled, so they pass an empty map and stay byte-identical. D1b — duplicate tool-call ids. Gemini returns the function name as the call id, so parallel calls to one function collide, and that id is what pairs a tool result back to its call. dedupe_provider_ids suffixes collisions (`get_weather`, `get_weather-2`); safe because both halves of the pairing are re-emitted from this same value. D2 — block-array content discarded (silent). parse_openai read `content` with `as_str()`, which yields "" for the typed-block arrays Gemini/Qwen35/gpt-oss send, so the model's answer vanished with no error. openai_content_parts now accepts a string or a block array, splitting text and reasoning (Gemini nests prose under `summary`); message-level reasoning_content/reasoning still win, so DeepSeek and vLLM hosts are unchanged. Also adds a turn-start log line to buzz-acp pool.rs, labelled to pair with log_stop_reason, so a stalled turn (start with no stop) is distinguishable from one that was never entered. Design write-up: docs/08-gemini-provider-fixes.md on the benchmark branch. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Atish Patel --- crates/buzz-acp/src/pool.rs | 27 +++- crates/buzz-agent/src/agent.rs | 3 + crates/buzz-agent/src/llm.rs | 227 +++++++++++++++++++++++++++++++-- crates/buzz-agent/src/types.rs | 13 +- 4 files changed, 254 insertions(+), 16 deletions(-) diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index b1fd68d044..038f8a714c 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -1897,6 +1897,18 @@ pub async fn run_prompt_task( None => prompt_sections.iter().map(String::as_str).collect(), }; + // Turn start, labelled exactly as `log_stop_reason` labels the end, so a + // log reads as start/stop pairs. Purely observational: an unpaired start is + // the only durable evidence that a turn was entered and never returned, and + // without it a stalled agent and an agent nobody woke leave identical logs — + // zero completions either way, so anything reading them afterwards has to + // guess which happened. + tracing::info!( + target: "pool::prompt", + "turn starting for {}", + prompt_label(&source) + ); + // When control_rx is Some (channel tasks), wrap the prompt in select! so // the main loop can cancel, interrupt, or rotate it. Heartbeats // (control_rx=None) take the simple await path — they are not controllable. @@ -3131,12 +3143,19 @@ fn classify_control_cancel_failure( } } -/// Log a stop reason at the appropriate tracing level. -fn log_stop_reason(source: &PromptSource, stop_reason: &StopReason) { - let label = match source { +/// How a turn's source is named in the `pool::prompt` log lines. +/// +/// Shared by the turn-start and turn-stop lines so a log can be read as pairs. +fn prompt_label(source: &PromptSource) -> String { + match source { PromptSource::Channel(cid) => format!("channel {cid}"), PromptSource::Heartbeat => "heartbeat".to_string(), - }; + } +} + +/// Log a stop reason at the appropriate tracing level. +fn log_stop_reason(source: &PromptSource, stop_reason: &StopReason) { + let label = prompt_label(source); match stop_reason { StopReason::EndTurn => { tracing::info!(target: "pool::prompt", "turn complete for {label}: end_turn"); diff --git a/crates/buzz-agent/src/agent.rs b/crates/buzz-agent/src/agent.rs index cbf27357f7..ed04daca2c 100644 --- a/crates/buzz-agent/src/agent.rs +++ b/crates/buzz-agent/src/agent.rs @@ -696,6 +696,9 @@ pub(crate) fn push_hook_outputs_as_tool_results( provider_id: provider_id.clone(), name: tool_name, arguments: serde_json::json!({}), + // Synthesised locally, so there is no provider wire form to + // preserve. + provider_extra: Default::default(), }], }); history.push(HistoryItem::ToolResult(ToolResult { diff --git a/crates/buzz-agent/src/llm.rs b/crates/buzz-agent/src/llm.rs index f69a963533..22d3f8b73e 100644 --- a/crates/buzz-agent/src/llm.rs +++ b/crates/buzz-agent/src/llm.rs @@ -3,7 +3,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use reqwest::Client; -use serde_json::{json, Value}; +use serde_json::{json, Map, Value}; use tokio::sync::Mutex; use tokio::time::Instant; @@ -848,11 +848,24 @@ fn openai_body( let calls: Vec = tool_calls .iter() .map(|c| { - json!({ - "id": c.provider_id, "type": "function", - "function": { "name": c.name, - "arguments": serde_json::to_string(&c.arguments) - .unwrap_or_else(|_| "{}".into()) } }) + let mut call = serde_json::Map::new(); + call.insert("id".into(), json!(c.provider_id)); + call.insert("type".into(), json!("function")); + // Provider-owned fields go back beside `function`, + // which is where the provider put them. Position is + // load-bearing, not cosmetic: Gemini rejects a + // `thoughtSignature` nested inside `function{}` with + // the same 400 it gives for one that is missing. + for (k, v) in &c.provider_extra { + call.insert(k.clone(), v.clone()); + } + call.insert( + "function".into(), + json!({ "name": c.name, + "arguments": serde_json::to_string(&c.arguments) + .unwrap_or_else(|_| "{}".into()) }), + ); + Value::Object(call) }) .collect(); msg.insert("tool_calls".into(), Value::Array(calls)); @@ -1120,10 +1133,15 @@ fn parse_responses(v: Value) -> Result { let args: Value = serde_json::from_str(raw).map_err(|e| { AgentError::Llm(format!("function_call.arguments not valid JSON: {e}")) })?; + // No passthrough on this route: `responses_body` replays a + // function call as `{call_id, name, arguments}` and the Responses + // API asks for nothing else, so an empty map keeps the request + // byte-identical to before. tool_calls.push(make_tool_call( str_field(item, "call_id"), str_field(item, "name"), args, + Default::default(), )?); } Some("reasoning") => { @@ -1301,6 +1319,76 @@ fn str_field(v: &Value, key: &str) -> String { v.get(key).and_then(Value::as_str).unwrap_or("").to_owned() } +/// Append `part` to `buf` on its own line, ignoring empties. +fn push_part(buf: &mut String, part: &str) { + if part.is_empty() { + return; + } + if !buf.is_empty() { + buf.push('\n'); + } + buf.push_str(part); +} + +/// Split an OpenAI-shaped `message.content` into `(text, reasoning)`. +/// +/// Standard OpenAI sends a string. Several models on the Databricks MLflow route +/// — Gemini, Qwen35, gpt-oss — send an array of typed blocks instead, and +/// `as_str()` yields nothing for an array, so their entire answer was being +/// discarded: no error, no warning, just a turn that looked like the model had +/// said nothing. `parse_anthropic` already walks a block array; this gives +/// `parse_openai` the same tolerance. +fn openai_content_parts(content: Option<&Value>) -> (String, String) { + let mut text = String::new(); + let mut reasoning = String::new(); + match content { + Some(Value::String(s)) => text.push_str(s), + Some(Value::Array(blocks)) => { + for b in blocks { + match b.get("type").and_then(Value::as_str) { + Some("text") => push_part(&mut text, &str_field(b, "text")), + Some("reasoning") => match b.get("summary").and_then(Value::as_array) { + // Gemini nests the prose one level down under `summary`. + Some(summary) => { + for s in summary { + push_part(&mut reasoning, &str_field(s, "text")); + } + } + None => push_part(&mut reasoning, &str_field(b, "text")), + }, + // An untyped block carrying text is still the model talking; + // treating it as text loses nothing and keeps one more + // provider out of the silent-empty-answer failure mode. + _ => push_part(&mut text, &str_field(b, "text")), + } + } + } + _ => {} + } + (text, reasoning) +} + +/// Make `provider_id` unique across one assistant turn's tool calls. +/// +/// Gemini returns the function name as the id, so two parallel calls to the same +/// function arrive sharing one id — and that id is what pairs a `role:"tool"` +/// result back to its call, leaving two results indistinguishable. Rewriting is +/// safe because both halves of that pairing are re-emitted from this same value; +/// the provider never sees its original id again. +fn dedupe_provider_ids(calls: &mut [ToolCall]) { + let mut seen: BTreeSet = BTreeSet::new(); + for c in calls.iter_mut() { + if seen.contains(&c.provider_id) { + let mut n = 2; + while seen.contains(&format!("{}-{n}", c.provider_id)) { + n += 1; + } + c.provider_id = format!("{}-{n}", c.provider_id); + } + seen.insert(c.provider_id.clone()); + } +} + fn parse_anthropic(v: Value) -> Result { let stop = map_stop(v.get("stop_reason").and_then(Value::as_str)); let mut tool_calls = Vec::new(); @@ -1323,10 +1411,12 @@ fn parse_anthropic(v: Value) -> Result { reasoning.push_str(t); } } + // Anthropic's replay shape is fully modelled, so nothing to keep. Some("tool_use") => tool_calls.push(make_tool_call( str_field(b, "id"), str_field(b, "name"), b.get("input").cloned().unwrap_or(Value::Null), + Default::default(), )?), _ => {} } @@ -1358,17 +1448,23 @@ fn parse_openai(v: Value) -> Result { let msg = choice .get("message") .ok_or_else(|| AgentError::Llm("missing message".into()))?; - let text = str_field(msg, "content"); + let (text, block_reasoning) = openai_content_parts(msg.get("content")); // DeepSeek and vLLM-style OpenAI-compat hosts expose reasoning tokens on the // message object. Prefer `reasoning_content` (DeepSeek's field name); fall - // back to `reasoning` (some other providers). Both are absent for standard - // OpenAI responses, which leaves this empty without any special-casing. + // back to `reasoning` (some other providers), and last to reasoning blocks + // found inside `content`. All three are absent for standard OpenAI + // responses, which leaves this empty without any special-casing. let reasoning = { let rc = str_field(msg, "reasoning_content"); - if rc.is_empty() { + let rc = if rc.is_empty() { str_field(msg, "reasoning") } else { rc + }; + if rc.is_empty() { + block_reasoning + } else { + rc } }; let mut tool_calls = Vec::new(); @@ -1380,13 +1476,25 @@ fn parse_openai(v: Value) -> Result { let raw = f.get("arguments").and_then(Value::as_str).unwrap_or("{}"); let args: Value = serde_json::from_str(raw) .map_err(|e| AgentError::Llm(format!("tool_call.arguments not valid JSON: {e}")))?; + // Everything on the wire object we do not model, kept for replay. + let extra = tc + .as_object() + .map(|o| { + o.iter() + .filter(|(k, _)| !matches!(k.as_str(), "id" | "type" | "function")) + .map(|(k, v)| (k.clone(), v.clone())) + .collect() + }) + .unwrap_or_default(); tool_calls.push(make_tool_call( str_field(tc, "id"), str_field(f, "name"), args, + extra, )?); } } + dedupe_provider_ids(&mut tool_calls); let input_tokens = openai_chat_input_tokens(&v); let output_tokens = sum_usage(&v, &["completion_tokens"]); let cached_input_tokens = openai_chat_cached_tokens(&v); @@ -1401,7 +1509,12 @@ fn parse_openai(v: Value) -> Result { }) } -fn make_tool_call(id: String, name: String, args: Value) -> Result { +fn make_tool_call( + id: String, + name: String, + args: Value, + provider_extra: Map, +) -> Result { if id.is_empty() || name.is_empty() { return Err(AgentError::Llm("tool_call missing id or name".into())); } @@ -1418,6 +1531,7 @@ fn make_tool_call(id: String, name: String, args: Value) -> Result Value { + json!({"choices": [{"finish_reason": "tool_calls", "message": { + "role": "assistant", + "content": [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "weighing it"}]}, + {"type": "text", "text": "391"} + ], + "tool_calls": [{ + "id": "get_weather", "type": "function", "thoughtSignature": "SIG-A", + "function": {"name": "get_weather", "arguments": "{\"city\":\"Paris\"}"} + }] + }}]}) + } + + #[test] + fn parse_openai_reads_text_out_of_a_block_array() { + // Before this, `as_str()` on the array yielded "" and the model's answer + // was discarded with no error at all. + let r = parse_openai(gemini_choice()).unwrap(); + assert_eq!(r.text, "391"); + assert_eq!(r.reasoning, "weighing it"); + } + + #[test] + fn parse_openai_still_reads_a_plain_string_content() { + let v = json!({"choices": [{"finish_reason": "stop", "message": { + "role": "assistant", "content": "plain"}}]}); + let r = parse_openai(v).unwrap(); + assert_eq!(r.text, "plain"); + assert_eq!(r.reasoning, ""); + } + + #[test] + fn parse_openai_keeps_unmodelled_tool_call_fields() { + let r = parse_openai(gemini_choice()).unwrap(); + let extra = &r.tool_calls[0].provider_extra; + assert_eq!(extra.get("thoughtSignature"), Some(&json!("SIG-A"))); + // `id`/`type`/`function` are modelled, so they must not be duplicated + // into the passthrough — they would be re-emitted twice. + assert!(!extra.contains_key("id")); + assert!(!extra.contains_key("type")); + assert!(!extra.contains_key("function")); + } + + #[test] + fn openai_body_replays_the_signature_beside_function_not_inside_it() { + // Position is what the gateway checks: nested inside `function{}` it is + // rejected with the same 400 as a missing signature. + let r = parse_openai(gemini_choice()).unwrap(); + let history = vec![HistoryItem::Assistant { + text: r.text.clone(), + tool_calls: r.tool_calls.clone(), + }]; + let body = openai_body( + &cfg(Provider::DatabricksV2), + "sys", + &history, + &[], + "databricks-gemini-3-6-flash", + None, + ); + let call = &body["messages"][1]["tool_calls"][0]; + assert_eq!(call["thoughtSignature"], json!("SIG-A")); + assert!(call["function"].get("thoughtSignature").is_none()); + assert_eq!(call["function"]["name"], json!("get_weather")); + } + + #[test] + fn parse_openai_makes_duplicate_tool_call_ids_unique() { + // Gemini returns the function name as the id, so parallel calls to one + // function collide and their results become indistinguishable. + let v = json!({"choices": [{"finish_reason": "tool_calls", "message": { + "role": "assistant", "content": "", + "tool_calls": [ + {"id": "get_weather", "type": "function", + "function": {"name": "get_weather", "arguments": "{\"city\":\"Paris\"}"}}, + {"id": "get_weather", "type": "function", + "function": {"name": "get_weather", "arguments": "{\"city\":\"Rome\"}"}} + ]}}]}); + let r = parse_openai(v).unwrap(); + assert_eq!(r.tool_calls[0].provider_id, "get_weather"); + assert_eq!(r.tool_calls[1].provider_id, "get_weather-2"); + } + #[test] fn parse_openai_uses_prompt_tokens() { let v = serde_json::json!({ diff --git a/crates/buzz-agent/src/types.rs b/crates/buzz-agent/src/types.rs index 1b5f30b1ce..c73bd879ad 100644 --- a/crates/buzz-agent/src/types.rs +++ b/crates/buzz-agent/src/types.rs @@ -1,5 +1,5 @@ use serde::Deserialize; -use serde_json::Value; +use serde_json::{Map, Value}; /// Byte-equivalent charged to the handoff/context-pressure gate for a single /// image tool result. The gate maps bytes to tokens at 1 byte/token (see @@ -108,6 +108,17 @@ pub struct ToolCall { pub provider_id: String, pub name: String, pub arguments: Value, + /// Fields the provider put on the tool call that we do not model, kept so + /// the assistant turn can be replayed the way it arrived. + /// + /// Gemini on the Databricks MLflow route returns a `thoughtSignature` per + /// call and *requires* it echoed back: replaying without it fails the whole + /// request with `Function call is missing a thought_signature in functionCall + /// parts`. For an agent loop that lands on the very first tool call, so the + /// model is unusable without this. Carrying whatever we did not model, + /// rather than naming that one field, means the next provider with an opaque + /// per-call token needs no change here. + pub provider_extra: Map, } #[derive(Debug, Clone)] From aaaaa792621748f6cc6eb6c542fef503f27c8d95 Mon Sep 17 00:00:00 2001 From: Atish Patel Date: Wed, 29 Jul 2026 10:54:49 -0500 Subject: [PATCH 2/2] fix(agent): count provider_extra in history size accounting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses codex review (P2). `HistoryItem::size_with` summed only provider_id + name + arguments, so a Gemini `thoughtSignature` — which is re-serialized into every replayed tool call — was invisible to `truncate_history` (request-body sizing) and the context-pressure/handoff gate. A large or repeated signature could push the real request past the configured byte/context budget before truncation kicked in. Count the serialized `provider_extra` in both measures. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Atish Patel --- crates/buzz-agent/src/types.rs | 40 ++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/crates/buzz-agent/src/types.rs b/crates/buzz-agent/src/types.rs index c73bd879ad..a3d48a7cf1 100644 --- a/crates/buzz-agent/src/types.rs +++ b/crates/buzz-agent/src/types.rs @@ -93,6 +93,13 @@ impl HistoryItem { + serde_json::to_vec(&c.arguments) .map(|b| b.len()) .unwrap_or(0) + // `provider_extra` (e.g. a Gemini + // `thoughtSignature`) is re-serialized into + // every replayed call, so it counts toward the + // request body and the context-pressure gate. + + serde_json::to_vec(&c.provider_extra) + .map(|b| b.len()) + .unwrap_or(0) }) .sum::() } @@ -364,6 +371,39 @@ mod tests { assert!(item.estimated_bytes() >= 3_118_884); } + #[test] + fn assistant_size_counts_provider_extra() { + // A Gemini `thoughtSignature` rides the wire on every replayed call, so + // both size measures must see it — otherwise `truncate_history` and the + // handoff gate under-count and let the real request exceed the budget. + let mut extra = Map::new(); + extra.insert("thoughtSignature".into(), Value::String("S".repeat(500))); + let with_extra = HistoryItem::Assistant { + text: String::new(), + tool_calls: vec![ToolCall { + provider_id: "id".into(), + name: "t".into(), + arguments: Value::Null, + provider_extra: extra, + }], + }; + let without_extra = HistoryItem::Assistant { + text: String::new(), + tool_calls: vec![ToolCall { + provider_id: "id".into(), + name: "t".into(), + arguments: Value::Null, + provider_extra: Map::new(), + }], + }; + assert!(with_extra.estimated_bytes() > without_extra.estimated_bytes() + 500); + assert_eq!( + with_extra.estimated_bytes(), + with_extra.context_pressure_bytes(), + "provider_extra is text, so both measures must agree" + ); + } + #[test] fn text_content_size_is_identical_for_both_measures() { // Only images diverge; text must size the same under both paths.