feat(agent): make Gemini and MLflow-route models usable through databricks_v2 (#3569)

## Summary

Makes Gemini — and every other non-Claude, non-GPT-5 model on the
Databricks MLflow route (`databricks_v2`) — usable in an agent loop.
These are the non-`benchmarks/` changes from
`benchmark/harness-accounting-and-solo`, lifted onto a clean base off
`main` so they can land independently while the harness work continues.

Two defects made these models unusable, one fatal and one silent. Both
live only in `openai_body` / `parse_openai`, which is the
least-exercised of the three `databricks_v2` sub-routes — the
`luna`/`sol` conditions run the Responses route and the `opus`
conditions run the Anthropic route, so **this change is inert for every
model already in use** and only lights up the MLflow path.

## Why a third route at all

`databricks_v2_route_for_model` buckets by model family: `claude*` →
Anthropic Messages, `gpt-5`/code-names → OpenAI Responses, **everything
else → MLflow chat-completions**. Gemini, Qwen, gpt-oss, and friends all
fall through to that third pair — and both bugs below live only there.

## D1 — dropped thought signatures (fatal)

Gemini returns a `thoughtSignature` on every tool call and **requires it
echoed back**. `openai_body` reserialized each call as `{id, type,
function}` only, dropping the field, so the next request 400'd:

```
HTTP 400  Function call is missing a thought_signature in functionCall parts.
```

For a coding agent this fires on the **first** tool call, so the model
never completes a single turn.

**Position is load-bearing.** A four-shape replay probe against the live
gateway established that the signature must sit as a *sibling* of
`function` — nesting it inside `function{}` fails with the *same* 400 as
omitting it. A fix that "preserves the field" without preserving its
position passes a unit test and still 400s.

The fix: `ToolCall` gains `provider_extra: Map<String, Value>`.
`parse_openai` captures every top-level wire key except the three we
model (`id`, `type`, `function`); `openai_body` re-emits them beside
`function`. Keeping *whatever we did not model*, rather than naming
`thoughtSignature`, means the next provider with an opaque per-call
token needs no change here. The Responses and Anthropic replay shapes
are fully modelled, so they pass `Default::default()` and stay
**byte-identical** to before.

### D1b — duplicate tool-call ids (same root cause)

Gemini returns the **function name** as the id, so two parallel calls to
one function arrive sharing an id — and that id is what pairs a
`role:"tool"` result back to its call, making two results
indistinguishable. `dedupe_provider_ids` suffixes collisions
(`get_weather`, `get_weather-2`). Safe because both halves of the
pairing (the assistant `tool_calls[].id` and the result's
`tool_call_id`) are re-emitted from this same value; the provider never
sees its original id again.

## D2 — block-array content discarded (silent, worse than a crash)

`parse_openai` read `content` with `as_str()`, which returns `""` for
anything that isn't a JSON string. Gemini (and Qwen35, gpt-oss) send an
array of typed blocks:

```json
"content": [
  {"type": "reasoning", "summary": [{"type": "summary_text", "text": "…"}]},
  {"type": "text", "text": "391"}
]
```

So the model answered and the answer was thrown away — no error, no
warning, just a turn that looked like the model had said nothing. On a
benchmark this reads as "Gemini is bad at the task" rather than "buzz
dropped the reply."

`openai_content_parts` now accepts either shape — string as before, or a
block array where `text` blocks concatenate into text and `reasoning`
blocks into reasoning (Gemini nests the prose one level down under
`summary`). Message-level `reasoning_content` / `reasoning` still win
when present, so DeepSeek and vLLM-style hosts are unchanged; block
reasoning is the last fallback.

## Also: a turn-start log line (`buzz-acp` `pool.rs`)

Small, independent observability change that also rides in the
non-benchmark delta: `run_prompt_task` now emits a `pool::prompt` "turn
starting" line, labelled by the same `prompt_label` helper as
`log_stop_reason`, so a log reads as start/stop pairs. An unpaired start
is the only durable evidence that a turn was entered and never returned
— without it, a stalled agent and an agent nobody woke leave identical
(zero-completion) logs.

## Interaction with #3538

#3538 (already merged) rewrote `databricks_v2_route_for_model` to route
by boundary-aware model-family segments. That change and this one touch
**different functions** in `llm.rs` — routing vs. body/parse — and
compose cleanly; the family routing decides *which* pair runs, and this
fixes the MLflow pair it can now select.

## Testing

- `cargo fmt --all -- --check`, `cargo clippy -p buzz-agent -p buzz-acp
--all-targets -- -D warnings` — clean.
- `cargo test -p buzz-agent -p buzz-acp` — all green (304 + 632 lib
tests plus integration suites, 0 failures). Five new tests cover:
block-array text extraction, plain-string regression, passthrough
capture (and non-duplication of the modelled keys), replay position
(`thoughtSignature` beside `function`, not inside it), and id
de-duplication.
- Wire evidence: the four-shape replay table and the reasoning-effort
probe were run against `block-lakehouse-staging` (recorded in the design
doc).

## Relationship to the benchmark branch

The full design write-up (four-shape replay table, position-matters
analysis, effort verification, and open pricing item) lives in
`docs/08-gemini-provider-fixes.md` on
`benchmark/harness-accounting-and-solo`. The benchmark manifests and
endpoint-config entries that exercise these models are separable and
stay on that branch.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Signed-off-by: Atish Patel <atish@squareup.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Atish Patel
2026-07-29 11:49:51 -05:00
committed by GitHub
parent 9beb3b8c6e
commit 4a1ebf25c7
4 changed files with 294 additions and 16 deletions

View File

@@ -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");

View File

@@ -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 {

View File

@@ -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<Value> = 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<LlmResponse, AgentError> {
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<String> = 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<LlmResponse, AgentError> {
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<LlmResponse, AgentError> {
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<LlmResponse, AgentError> {
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<LlmResponse, AgentError> {
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<LlmResponse, AgentError> {
})
}
fn make_tool_call(id: String, name: String, args: Value) -> Result<ToolCall, AgentError> {
fn make_tool_call(
id: String,
name: String,
args: Value,
provider_extra: Map<String, Value>,
) -> Result<ToolCall, AgentError> {
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<ToolCall, Age
provider_id: id,
name,
arguments,
provider_extra,
})
}
@@ -2365,6 +2479,7 @@ mod tests {
provider_id: "toolu_1".into(),
name: "dev__view_image".into(),
arguments: serde_json::json!({"source":"x.png"}),
provider_extra: Default::default(),
}],
},
HistoryItem::ToolResult(ToolResult {
@@ -2414,6 +2529,7 @@ mod tests {
provider_id: "call_abc".into(),
name: "dev__shell".into(),
arguments: serde_json::json!({"command": "ls"}),
provider_extra: Default::default(),
}],
},
HistoryItem::ToolResult(ToolResult {
@@ -2512,6 +2628,7 @@ mod tests {
provider_id: "call_x".into(),
name: "t".into(),
arguments: serde_json::json!({}),
provider_extra: Default::default(),
}],
},
];
@@ -2790,11 +2907,13 @@ mod tests {
provider_id: "toolu_a".into(),
name: "dev__view_image".into(),
arguments: serde_json::json!({"source": "a.png"}),
provider_extra: Default::default(),
},
ToolCall {
provider_id: "toolu_b".into(),
name: "dev__view_image".into(),
arguments: serde_json::json!({"source": "b.png"}),
provider_extra: Default::default(),
},
],
},
@@ -3819,6 +3938,92 @@ mod tests {
assert_eq!(parse_anthropic(v).unwrap().input_tokens, None);
}
/// A Gemini reply as the Databricks MLflow route actually returns it:
/// block-array `content`, and a `thoughtSignature` beside `function`.
fn gemini_choice() -> 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!({

View File

@@ -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
@@ -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::<usize>()
}
@@ -108,6 +115,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<String, Value>,
}
#[derive(Debug, Clone)]
@@ -353,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.