mirror of
https://github.com/block/buzz.git
synced 2026-11-02 15:31:27 +00:00
16d4ec335e
Uses MeshLLM built-in `mesh` collective intelligence / Mixture of Agents when Buzz Auto sees two or more distinct physical models. With zero or one model, Auto remains ordinary `auto`. This is an alternative way to improve tool responses, accuracy, and resistance to hallucination when a high-latency distributed mesh contains diverse models. Models and members may come and go: collective routing enables only after stable capacity, drops on confirmed contraction, and can recover later. Mesh-specific failures retry once through ordinary Auto. This update also pins MeshLLM to a v0.73.1-compatible backport of [MeshLLM #1074](https://github.com/Mesh-LLM/mesh-llm/pull/1074), so client-only Buzz nodes cannot enter model election or download a remote provider model. Buzz preserves the selected local sharing model and switches an existing client to sharing across a controlled app restart, retaining one runtime and one `:9337` / `:3131` pair per machine. Validation: - Full local `just ci` passes on the cleaned branch. - MeshLLM host-runtime suite: 1,568 passed, 0 failed; strict Clippy passes. - Buzz desktop Tauri suite with `mesh-llm`: 1,721 passed, 0 failed; strict feature Clippy passes. - Playwright covers client-to-share using the saved local model and no destructive stop. - Packaged two-machine testing proved single-model routing, dual-model collective routing, tool-markup fallback, and runtime reuse. - Packaged client-only recheck routed a real Mini Buzz turn through M5 while Mini stayed `is_client=true`, `is_host=false`, hosted no models, and created no Gemma cache. Builds on the recovery work merged in #2823; this PR does not duplicate it. --------- Signed-off-by: Michael Neale <michael.neale@gmail.com> Signed-off-by: Tyler Longwell <tlongwell@block.xyz> Co-authored-by: Michael Neale <michael.neale@gmail.com> Co-authored-by: npub1qyvc0c5kl4gqv2fd97fsk46tu378sqgy35vc83rvgfwne90sel7s0ed67d <011987e296fd5006292d2f930b574be47c7801048d1983c46c425d3c95f0cffd@buzz.block.builderlab.xyz> Co-authored-by: Tyler Longwell <tlongwell@block.xyz>
3847 lines
144 KiB
Rust
3847 lines
144 KiB
Rust
use std::collections::BTreeSet;
|
||
use std::sync::atomic::{AtomicBool, Ordering};
|
||
use std::sync::Arc;
|
||
|
||
use reqwest::Client;
|
||
use serde_json::{json, Value};
|
||
use tokio::sync::Mutex;
|
||
use tokio::time::Instant;
|
||
|
||
use crate::auth::{PkceOAuthConfig, PkceOAuthTokenSource, StaticTokenSource, TokenSource};
|
||
use crate::config::{
|
||
is_openai_host, normalize_effort_for_anthropic_route, normalize_effort_for_openai_route,
|
||
Config, OpenAiApi, Provider, ThinkingEffort,
|
||
};
|
||
use crate::types::{
|
||
AgentError, HistoryItem, LlmResponse, ProviderStop, ToolCall, ToolDef, ToolResultContent,
|
||
};
|
||
|
||
/// Databricks OAuth client_id — the public Databricks-published CLI client.
|
||
/// PKCE-only, no secret. Same identifier goose uses, so a user's browser
|
||
/// consent for `databricks-cli` covers buzz-agent too.
|
||
const DATABRICKS_CLIENT_ID: &str = "databricks-cli";
|
||
const DATABRICKS_OAUTH_SCOPES: &[&str] = &["all-apis", "offline_access"];
|
||
|
||
const MAX_LLM_RESPONSE_BYTES: usize = 16 * 1024 * 1024;
|
||
const MAX_LLM_ERROR_BODY_BYTES: usize = 4 * 1024;
|
||
const STALL_NOTICE_THRESHOLD: std::time::Duration = std::time::Duration::from_secs(300);
|
||
const MESH_VIRTUAL_MODEL_ID: &str = "mesh";
|
||
const MESH_AUTO_MODEL_ID: &str = "auto";
|
||
const MESH_AUTO_CATALOG_TTL: std::time::Duration = std::time::Duration::from_secs(5);
|
||
const MESH_AUTO_CATALOG_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(2);
|
||
const MESH_AUTO_COOLDOWN: std::time::Duration = std::time::Duration::from_secs(30);
|
||
const MESH_AUTO_ENABLE_OBSERVATIONS: u8 = 2;
|
||
const MESH_MOA_UNAVAILABLE_MESSAGE: &str = "MoA requires ≥2 models available in the mesh";
|
||
|
||
/// Parser for an OpenAI-family JSON response. Per-endpoint pair lives
|
||
/// alongside its `_body` serializer.
|
||
type OpenAiParse = fn(Value) -> Result<LlmResponse, AgentError>;
|
||
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
enum MeshCatalogObservation {
|
||
Available,
|
||
Unavailable,
|
||
Unknown,
|
||
}
|
||
|
||
fn mesh_catalog_supports_collective(catalog: &Value) -> Option<bool> {
|
||
let models = catalog.get("data").and_then(Value::as_array)?;
|
||
let mut has_virtual_mesh = false;
|
||
let mut physical_models = BTreeSet::new();
|
||
for id in models
|
||
.iter()
|
||
.filter_map(|model| model.get("id").and_then(Value::as_str))
|
||
.map(str::trim)
|
||
.filter(|id| !id.is_empty())
|
||
{
|
||
if id == MESH_VIRTUAL_MODEL_ID {
|
||
has_virtual_mesh = true;
|
||
} else if id != MESH_AUTO_MODEL_ID {
|
||
physical_models.insert(id.replace("@main", ""));
|
||
}
|
||
}
|
||
Some(has_virtual_mesh && physical_models.len() >= 2)
|
||
}
|
||
|
||
fn looks_like_unstructured_tool_call(text: &str) -> bool {
|
||
let text = text.trim_start().to_ascii_lowercase();
|
||
text.starts_with("<|tool_call")
|
||
|| text.starts_with("<tool_call")
|
||
|| text.starts_with("[tool_call]")
|
||
}
|
||
|
||
#[derive(Debug, Default)]
|
||
struct MeshAutoState {
|
||
last_checked: Option<Instant>,
|
||
consecutive_available: u8,
|
||
collective_enabled: bool,
|
||
cooldown_until: Option<Instant>,
|
||
}
|
||
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
enum DatabricksV2Route {
|
||
OpenAiResponses,
|
||
AnthropicMessages,
|
||
MlflowChatCompletions,
|
||
}
|
||
|
||
pub struct Llm {
|
||
http: Client,
|
||
/// One-shot sticky flag: set when a Chat Completions request comes
|
||
/// back with a "use /v1/responses" provider error while `cfg.openai_api
|
||
/// == Auto`. Subsequent OpenAI calls then go straight to Responses
|
||
/// for the lifetime of the process.
|
||
auto_upgraded: AtomicBool,
|
||
/// Hysteretic view of whether mesh-llm currently advertises its virtual
|
||
/// Mixture-of-Agents model. A TTL, confirmation count, and failure cooldown
|
||
/// let long-running agents adapt without bouncing as peers briefly flap.
|
||
mesh_auto_state: Mutex<MeshAutoState>,
|
||
/// Bearer-token source for OpenAI-family requests. Static for OpenAI
|
||
/// (the `OPENAI_COMPAT_API_KEY` env var) and Databricks-with-token
|
||
/// (the `DATABRICKS_TOKEN` env var); a refreshable PKCE engine for
|
||
/// Databricks otherwise. Anthropic doesn't use this — it always
|
||
/// reads `cfg.api_key` directly because the API expects `x-api-key`.
|
||
auth: Arc<dyn TokenSource>,
|
||
}
|
||
|
||
impl Llm {
|
||
pub fn new(cfg: &Config) -> Result<Self, AgentError> {
|
||
let http = Client::builder()
|
||
.connect_timeout(std::time::Duration::from_secs(10))
|
||
.read_timeout(cfg.llm_timeout)
|
||
.build()
|
||
.map_err(|e| AgentError::Llm(format!("http: {e}")))?;
|
||
let auth = build_token_source(cfg)?;
|
||
Ok(Self {
|
||
http,
|
||
auto_upgraded: AtomicBool::new(false),
|
||
mesh_auto_state: Mutex::new(MeshAutoState::default()),
|
||
auth,
|
||
})
|
||
}
|
||
|
||
pub async fn complete(
|
||
&self,
|
||
cfg: &Config,
|
||
system_prompt: &str,
|
||
history: &[HistoryItem],
|
||
tools: &[ToolDef],
|
||
effective_model: &str,
|
||
) -> Result<LlmResponse, AgentError> {
|
||
let effort = cfg.thinking_effort;
|
||
let result = match cfg.provider {
|
||
Provider::Anthropic => {
|
||
let v = self
|
||
.post_anthropic(
|
||
cfg,
|
||
&anthropic_body(
|
||
cfg,
|
||
system_prompt,
|
||
history,
|
||
tools,
|
||
effective_model,
|
||
effort,
|
||
),
|
||
)
|
||
.await?;
|
||
parse_anthropic(v)
|
||
}
|
||
Provider::OpenAi | Provider::Databricks => {
|
||
self.openai_request(
|
||
cfg,
|
||
effective_model,
|
||
!tools.is_empty(),
|
||
|use_responses, request_model| {
|
||
// Normalize effort for model-specific availability. Startup no longer rejects
|
||
// `max` for pure OpenAI/Databricks; this per-model table is the single authority
|
||
// — it keeps `max` for gpt-5.6, clamps `max`→`xhigh` for other OpenAI-shaped
|
||
// models, and still applies corrections like none→minimal on the gpt-5 base.
|
||
let e =
|
||
effort.map(|ef| normalize_effort_for_openai_route(ef, request_model));
|
||
if use_responses {
|
||
(
|
||
responses_body(
|
||
cfg,
|
||
system_prompt,
|
||
history,
|
||
tools,
|
||
request_model,
|
||
e,
|
||
),
|
||
parse_responses as OpenAiParse,
|
||
)
|
||
} else {
|
||
(
|
||
openai_body(cfg, system_prompt, history, tools, request_model, e),
|
||
parse_openai as OpenAiParse,
|
||
)
|
||
}
|
||
},
|
||
)
|
||
.await
|
||
}
|
||
Provider::DatabricksV2 => {
|
||
self.databricks_v2_request(cfg, effective_model, |route| match route {
|
||
DatabricksV2Route::OpenAiResponses => {
|
||
// OpenAI Responses path: normalize effort against the per-model table.
|
||
let e =
|
||
effort.map(|ef| normalize_effort_for_openai_route(ef, effective_model));
|
||
(
|
||
responses_body(cfg, system_prompt, history, tools, effective_model, e),
|
||
parse_responses as OpenAiParse,
|
||
)
|
||
}
|
||
DatabricksV2Route::AnthropicMessages => {
|
||
// Anthropic Messages path: normalize effort (none|minimal → omit).
|
||
let e = effort.and_then(normalize_effort_for_anthropic_route);
|
||
(
|
||
anthropic_body(cfg, system_prompt, history, tools, effective_model, e),
|
||
parse_anthropic as OpenAiParse,
|
||
)
|
||
}
|
||
DatabricksV2Route::MlflowChatCompletions => {
|
||
// MLflow Chat path (OpenAI-shaped): normalize effort against the per-model table.
|
||
let e =
|
||
effort.map(|ef| normalize_effort_for_openai_route(ef, effective_model));
|
||
(
|
||
openai_body(cfg, system_prompt, history, tools, effective_model, e),
|
||
parse_openai as OpenAiParse,
|
||
)
|
||
}
|
||
})
|
||
.await
|
||
}
|
||
};
|
||
// Stamp the effective model into Llm errors so log lines carry
|
||
// `llm: (model-name) 404 Not Found: …` instead of the bare status.
|
||
// The `llm: ` prefix comes from `Display for AgentError::Llm`; the
|
||
// map_err here prepends `(model-name) ` to the inner string only.
|
||
// This is the single place all provider paths converge, so the mapping
|
||
// is centralized and never needs to be repeated in each provider arm.
|
||
result.map_err(|e| match e {
|
||
AgentError::Llm(s) => AgentError::Llm(format!("({effective_model}) {s}")),
|
||
AgentError::LlmModelNotFound(s) => {
|
||
AgentError::LlmModelNotFound(format!("({effective_model}) {s}"))
|
||
}
|
||
other => other,
|
||
})
|
||
}
|
||
|
||
pub async fn summarize(
|
||
&self,
|
||
cfg: &Config,
|
||
system_prompt: &str,
|
||
user_prompt: &str,
|
||
max_output_tokens: u32,
|
||
effective_model: &str,
|
||
) -> Result<String, AgentError> {
|
||
match cfg.provider {
|
||
Provider::Anthropic => {
|
||
let body = json!({
|
||
"model": effective_model,
|
||
"max_tokens": max_output_tokens,
|
||
"system": system_prompt,
|
||
"messages": [{
|
||
"role": "user",
|
||
"content": [{ "type": "text", "text": user_prompt }],
|
||
}],
|
||
});
|
||
Ok(parse_anthropic(self.post_anthropic(cfg, &body).await?)?.text)
|
||
}
|
||
Provider::OpenAi | Provider::Databricks => {
|
||
let r = self
|
||
.openai_request(
|
||
cfg,
|
||
effective_model,
|
||
false,
|
||
|use_responses, request_model| {
|
||
if use_responses {
|
||
(
|
||
json!({
|
||
"model": request_model,
|
||
"max_output_tokens": max_output_tokens,
|
||
"instructions": system_prompt,
|
||
"input": user_prompt,
|
||
}),
|
||
parse_responses as OpenAiParse,
|
||
)
|
||
} else {
|
||
(
|
||
json!({
|
||
"model": request_model,
|
||
"stream": false,
|
||
"max_completion_tokens": max_output_tokens,
|
||
"messages": [
|
||
{ "role": "system", "content": system_prompt },
|
||
{ "role": "user", "content": user_prompt },
|
||
],
|
||
}),
|
||
parse_openai as OpenAiParse,
|
||
)
|
||
}
|
||
},
|
||
)
|
||
.await?;
|
||
Ok(r.text)
|
||
}
|
||
Provider::DatabricksV2 => {
|
||
let r = self
|
||
.databricks_v2_request(cfg, effective_model, |route| match route {
|
||
DatabricksV2Route::OpenAiResponses => (
|
||
json!({
|
||
"model": effective_model,
|
||
"max_output_tokens": max_output_tokens,
|
||
"instructions": system_prompt,
|
||
"input": user_prompt,
|
||
}),
|
||
parse_responses as OpenAiParse,
|
||
),
|
||
DatabricksV2Route::AnthropicMessages => (
|
||
json!({
|
||
"model": effective_model,
|
||
"max_tokens": max_output_tokens,
|
||
"system": system_prompt,
|
||
"messages": [{
|
||
"role": "user",
|
||
"content": [{ "type": "text", "text": user_prompt }],
|
||
}],
|
||
}),
|
||
parse_anthropic as OpenAiParse,
|
||
),
|
||
DatabricksV2Route::MlflowChatCompletions => (
|
||
json!({
|
||
"model": effective_model,
|
||
"stream": false,
|
||
"max_completion_tokens": max_output_tokens,
|
||
"messages": [
|
||
{ "role": "system", "content": system_prompt },
|
||
{ "role": "user", "content": user_prompt },
|
||
],
|
||
}),
|
||
parse_openai as OpenAiParse,
|
||
),
|
||
})
|
||
.await?;
|
||
Ok(r.text)
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn post_anthropic(&self, cfg: &Config, body: &Value) -> Result<Value, AgentError> {
|
||
let url = format!("{}/v1/messages", cfg.base_url.trim_end_matches('/'));
|
||
post(&self.http, &url, body, false, |r| {
|
||
r.header("x-api-key", &cfg.api_key)
|
||
.header("anthropic-version", &cfg.anthropic_api_version)
|
||
})
|
||
.await
|
||
.map_err(PostError::into_agent)
|
||
}
|
||
|
||
/// OpenAI dispatch with Buzz's relay-mesh `auto` policy layered over the
|
||
/// normal endpoint selection. When enabled, a live virtual `mesh` model is
|
||
/// preferred; if the mesh contracts between discovery and inference, retry
|
||
/// the same request once through the router's ordinary `auto` model.
|
||
async fn openai_request<F>(
|
||
&self,
|
||
cfg: &Config,
|
||
effective_model: &str,
|
||
tools_supplied: bool,
|
||
mut build: F,
|
||
) -> Result<LlmResponse, AgentError>
|
||
where
|
||
F: FnMut(bool, &str) -> (Value, OpenAiParse) + Send,
|
||
{
|
||
let request_model = self.resolve_openai_model(cfg, effective_model).await;
|
||
let adaptive_mesh =
|
||
effective_model == MESH_AUTO_MODEL_ID && request_model == MESH_VIRTUAL_MODEL_ID;
|
||
|
||
let first = self
|
||
.openai_request_for_model(cfg, &request_model, &mut build)
|
||
.await;
|
||
match first {
|
||
Err(PostError::MeshFallback(detail)) if adaptive_mesh => {
|
||
self.cool_down_collective().await;
|
||
tracing::warn!(
|
||
configured_model = effective_model,
|
||
attempted_model = MESH_VIRTUAL_MODEL_ID,
|
||
fallback_model = MESH_AUTO_MODEL_ID,
|
||
provider_message = detail,
|
||
"relay-mesh auto: collective request failed; retrying once with auto"
|
||
);
|
||
self.openai_request_for_model(cfg, MESH_AUTO_MODEL_ID, &mut build)
|
||
.await
|
||
.map_err(PostError::into_agent)
|
||
}
|
||
Ok(response)
|
||
if adaptive_mesh
|
||
&& tools_supplied
|
||
&& response.tool_calls.is_empty()
|
||
&& looks_like_unstructured_tool_call(&response.text) =>
|
||
{
|
||
self.cool_down_collective().await;
|
||
tracing::warn!(
|
||
configured_model = effective_model,
|
||
attempted_model = MESH_VIRTUAL_MODEL_ID,
|
||
fallback_model = MESH_AUTO_MODEL_ID,
|
||
"relay-mesh auto: collective response emitted unstructured tool markup; retrying once with auto"
|
||
);
|
||
self.openai_request_for_model(cfg, MESH_AUTO_MODEL_ID, &mut build)
|
||
.await
|
||
.map_err(PostError::into_agent)
|
||
}
|
||
Ok(response) => Ok(response),
|
||
Err(error) => Err(error.into_agent()),
|
||
}
|
||
}
|
||
|
||
async fn cool_down_collective(&self) {
|
||
let now = Instant::now();
|
||
let mut state = self.mesh_auto_state.lock().await;
|
||
state.last_checked = Some(now);
|
||
state.consecutive_available = 0;
|
||
state.collective_enabled = false;
|
||
state.cooldown_until = Some(now + MESH_AUTO_COOLDOWN);
|
||
}
|
||
|
||
/// Resolve the model for one OpenAI-family request. Explicit model choices
|
||
/// are never changed. Relay-mesh `auto` dynamically follows the short-lived
|
||
/// `/models` catalog so long-running agents can adopt or leave MoA without
|
||
/// being restarted.
|
||
async fn resolve_openai_model(&self, cfg: &Config, effective_model: &str) -> String {
|
||
if cfg.provider != Provider::OpenAi
|
||
|| !cfg.prefer_mesh_for_auto
|
||
|| effective_model != MESH_AUTO_MODEL_ID
|
||
{
|
||
return effective_model.to_string();
|
||
}
|
||
|
||
let mut state = self.mesh_auto_state.lock().await;
|
||
let now = Instant::now();
|
||
if let Some(cooldown_until) = state.cooldown_until {
|
||
if now < cooldown_until {
|
||
return MESH_AUTO_MODEL_ID.to_string();
|
||
}
|
||
state.cooldown_until = None;
|
||
}
|
||
if state
|
||
.last_checked
|
||
.is_some_and(|checked_at| checked_at.elapsed() < MESH_AUTO_CATALOG_TTL)
|
||
{
|
||
return if state.collective_enabled {
|
||
MESH_VIRTUAL_MODEL_ID.to_string()
|
||
} else {
|
||
MESH_AUTO_MODEL_ID.to_string()
|
||
};
|
||
}
|
||
|
||
let observation = self.observe_mesh_virtual_model(cfg).await;
|
||
let checked_at = Instant::now();
|
||
state.last_checked = Some(checked_at);
|
||
match observation {
|
||
MeshCatalogObservation::Available => {
|
||
state.consecutive_available = state.consecutive_available.saturating_add(1);
|
||
if state.consecutive_available >= MESH_AUTO_ENABLE_OBSERVATIONS {
|
||
state.collective_enabled = true;
|
||
}
|
||
}
|
||
MeshCatalogObservation::Unavailable => {
|
||
if state.collective_enabled {
|
||
state.cooldown_until = Some(checked_at + MESH_AUTO_COOLDOWN);
|
||
}
|
||
state.consecutive_available = 0;
|
||
state.collective_enabled = false;
|
||
}
|
||
// A failed catalog check is not evidence that a previously stable
|
||
// mesh disappeared. Preserve the last confirmed state; inference
|
||
// itself remains authoritative and has the narrow 503 fallback.
|
||
MeshCatalogObservation::Unknown => {}
|
||
}
|
||
let request_model = if state.collective_enabled {
|
||
MESH_VIRTUAL_MODEL_ID
|
||
} else {
|
||
MESH_AUTO_MODEL_ID
|
||
};
|
||
tracing::debug!(
|
||
configured_model = effective_model,
|
||
request_model,
|
||
"relay-mesh auto: resolved request model from live catalog"
|
||
);
|
||
request_model.to_string()
|
||
}
|
||
|
||
async fn observe_mesh_virtual_model(&self, cfg: &Config) -> MeshCatalogObservation {
|
||
let url = format!("{}/models", cfg.base_url.trim_end_matches('/'));
|
||
let bearer = match self.auth.bearer().await {
|
||
Ok(bearer) => bearer,
|
||
Err(error) => {
|
||
tracing::debug!(
|
||
%error,
|
||
"relay-mesh auto: catalog auth unavailable; preserving last confirmed route"
|
||
);
|
||
return MeshCatalogObservation::Unknown;
|
||
}
|
||
};
|
||
let response = match self
|
||
.http
|
||
.get(&url)
|
||
.bearer_auth(bearer)
|
||
.timeout(MESH_AUTO_CATALOG_TIMEOUT)
|
||
.send()
|
||
.await
|
||
{
|
||
Ok(response) => response,
|
||
Err(error) => {
|
||
tracing::debug!(
|
||
%error,
|
||
"relay-mesh auto: catalog probe failed; preserving last confirmed route"
|
||
);
|
||
return MeshCatalogObservation::Unknown;
|
||
}
|
||
};
|
||
if !response.status().is_success() {
|
||
tracing::debug!(
|
||
status = %response.status(),
|
||
"relay-mesh auto: catalog probe was not successful; preserving last confirmed route"
|
||
);
|
||
return MeshCatalogObservation::Unknown;
|
||
}
|
||
match response.json::<Value>().await {
|
||
Ok(catalog) => {
|
||
let Some(available) = mesh_catalog_supports_collective(&catalog) else {
|
||
tracing::debug!(
|
||
"relay-mesh auto: catalog response has no data array; preserving last confirmed route"
|
||
);
|
||
return MeshCatalogObservation::Unknown;
|
||
};
|
||
if available {
|
||
MeshCatalogObservation::Available
|
||
} else {
|
||
MeshCatalogObservation::Unavailable
|
||
}
|
||
}
|
||
Err(error) => {
|
||
tracing::debug!(
|
||
%error,
|
||
"relay-mesh auto: invalid catalog response; preserving last confirmed route"
|
||
);
|
||
MeshCatalogObservation::Unknown
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Dispatch one OpenAI request for an already-resolved model. Endpoint
|
||
/// selection remains unchanged: pinned > sticky-upgraded > host default,
|
||
/// with the existing one-shot Chat-to-Responses upgrade.
|
||
async fn openai_request_for_model<F>(
|
||
&self,
|
||
cfg: &Config,
|
||
request_model: &str,
|
||
build: &mut F,
|
||
) -> Result<LlmResponse, PostError>
|
||
where
|
||
F: FnMut(bool, &str) -> (Value, OpenAiParse) + Send,
|
||
{
|
||
let use_responses = self.auto_upgraded.load(Ordering::Relaxed)
|
||
|| matches!(cfg.openai_api, OpenAiApi::Responses)
|
||
|| matches!(cfg.openai_api, OpenAiApi::Auto) && is_openai_host(&cfg.base_url);
|
||
|
||
if use_responses {
|
||
let (body, parse) = build(true, request_model);
|
||
return parse(
|
||
self.post_openai(cfg, "/responses", &body, request_model)
|
||
.await?,
|
||
)
|
||
.map_err(PostError::from);
|
||
}
|
||
let (body, parse) = build(false, request_model);
|
||
match self
|
||
.post_openai(cfg, "/chat/completions", &body, request_model)
|
||
.await
|
||
{
|
||
Ok(value) => parse(value).map_err(PostError::from),
|
||
Err(PostError::Agent(error))
|
||
if cfg.openai_api == OpenAiApi::Auto && self.try_upgrade(&error) =>
|
||
{
|
||
let (body, parse) = build(true, request_model);
|
||
parse(
|
||
self.post_openai(cfg, "/responses", &body, request_model)
|
||
.await?,
|
||
)
|
||
.map_err(PostError::from)
|
||
}
|
||
Err(error) => Err(error),
|
||
}
|
||
}
|
||
|
||
async fn databricks_v2_request<F>(
|
||
&self,
|
||
cfg: &Config,
|
||
effective_model: &str,
|
||
build: F,
|
||
) -> Result<LlmResponse, AgentError>
|
||
where
|
||
F: FnOnce(DatabricksV2Route) -> (Value, OpenAiParse) + Send,
|
||
{
|
||
let route = databricks_v2_route_for_model(effective_model);
|
||
let (body, parse) = build(route);
|
||
parse(
|
||
self.post_openai(cfg, databricks_v2_path(route), &body, effective_model)
|
||
.await
|
||
.map_err(PostError::into_agent)?,
|
||
)
|
||
}
|
||
|
||
/// POST to an OpenAI-family endpoint. For OpenAI-compat this is just
|
||
/// `{base_url}{path}` with the body untouched. For Databricks the URL
|
||
/// becomes `{base_url}/serving-endpoints/{model}/invocations` and the
|
||
/// `model` field is stripped from the body (Databricks rejects it —
|
||
/// the endpoint path already names the model).
|
||
async fn post_openai(
|
||
&self,
|
||
cfg: &Config,
|
||
path: &str,
|
||
body: &Value,
|
||
effective_model: &str,
|
||
) -> Result<Value, PostError> {
|
||
let (url, body_owned);
|
||
let body_ref: &Value = match cfg.provider {
|
||
Provider::Databricks => {
|
||
url = format!(
|
||
"{}/serving-endpoints/{}/invocations",
|
||
cfg.base_url.trim_end_matches('/'),
|
||
effective_model
|
||
);
|
||
body_owned = strip_model(body);
|
||
&body_owned
|
||
}
|
||
_ => {
|
||
url = format!("{}{}", cfg.base_url.trim_end_matches('/'), path);
|
||
body
|
||
}
|
||
};
|
||
|
||
// A 401 or 403 can mean the local expiry clock disagreed with the
|
||
// server (skew, revocation, a node that never saw the token). On the
|
||
// first such rejection, force a refresh keyed off the rejected bearer
|
||
// and retry once. The guard is local to this call so an earlier turn's
|
||
// rejection can never suppress a later turn's legitimate retry. Both
|
||
// statuses map to `LlmAuth` in `post`: a 403 is indistinguishable from
|
||
// an expired-token 403 here, so we refresh once and let it propagate.
|
||
let mut bearer = self.auth.bearer().await.map_err(PostError::from)?;
|
||
let mut refreshed = false;
|
||
loop {
|
||
match post(
|
||
&self.http,
|
||
&url,
|
||
body_ref,
|
||
effective_model == MESH_VIRTUAL_MODEL_ID,
|
||
|r| r.bearer_auth(&bearer),
|
||
)
|
||
.await
|
||
{
|
||
Err(PostError::Agent(AgentError::LlmAuth(_))) if !refreshed => {
|
||
refreshed = true;
|
||
bearer = self
|
||
.auth
|
||
.refresh_now(&bearer)
|
||
.await
|
||
.map_err(PostError::from)?;
|
||
}
|
||
result => return result,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// If `err` names `/v1/responses` / "use the Responses API", latch a
|
||
/// sticky upgrade so subsequent OpenAI calls hit Responses. Logged once.
|
||
fn try_upgrade(&self, err: &AgentError) -> bool {
|
||
let body = match err {
|
||
AgentError::Llm(s) => s.as_str(),
|
||
_ => return false, // auth/transport aren't "use the other endpoint" signals
|
||
};
|
||
if !is_responses_required_error(body) {
|
||
return false;
|
||
}
|
||
if !self.auto_upgraded.swap(true, Ordering::Relaxed) {
|
||
tracing::warn!(
|
||
provider_message = body,
|
||
"openai: provider asked for the Responses API; \
|
||
routing subsequent OpenAI calls to /v1/responses for this process"
|
||
);
|
||
}
|
||
true
|
||
}
|
||
}
|
||
|
||
fn anthropic_body(
|
||
cfg: &Config,
|
||
system_prompt: &str,
|
||
history: &[HistoryItem],
|
||
tools: &[ToolDef],
|
||
effective_model: &str,
|
||
effort: Option<ThinkingEffort>,
|
||
) -> Value {
|
||
let mut messages: Vec<Value> = Vec::new();
|
||
let mut pending: Vec<Value> = Vec::new();
|
||
let flush = |out: &mut Vec<Value>, p: &mut Vec<Value>| {
|
||
if !p.is_empty() {
|
||
out.push(json!({ "role": "user", "content": std::mem::take(p) }));
|
||
}
|
||
};
|
||
for item in history {
|
||
match item {
|
||
HistoryItem::User(text) => {
|
||
flush(&mut messages, &mut pending);
|
||
messages.push(json!({ "role": "user",
|
||
"content": [{ "type": "text", "text": text }] }));
|
||
}
|
||
HistoryItem::Assistant { text, tool_calls } => {
|
||
flush(&mut messages, &mut pending);
|
||
let mut content: Vec<Value> = Vec::new();
|
||
if !text.is_empty() {
|
||
content.push(json!({ "type": "text", "text": text }));
|
||
}
|
||
for c in tool_calls {
|
||
content.push(json!({ "type": "tool_use", "id": c.provider_id,
|
||
"name": c.name, "input": c.arguments }));
|
||
}
|
||
if content.is_empty() {
|
||
// Empty assistant turn (no text, no tool calls) — skip it.
|
||
// Anthropic rejects empty text blocks, and a placeholder
|
||
// just defers the problem. No tool_use = no pairing
|
||
// constraint, so omitting is safe.
|
||
continue;
|
||
}
|
||
messages.push(json!({ "role": "assistant", "content": content }));
|
||
}
|
||
HistoryItem::ToolResult(r) => pending.push(json!({
|
||
"type": "tool_result", "tool_use_id": r.provider_id,
|
||
"content": anthropic_tool_result_content(&r.content), "is_error": r.is_error })),
|
||
}
|
||
}
|
||
flush(&mut messages, &mut pending);
|
||
let tools_json: Vec<Value> = tools
|
||
.iter()
|
||
.map(|t| {
|
||
json!({
|
||
"name": t.name, "description": t.description, "input_schema": t.input_schema })
|
||
})
|
||
.collect();
|
||
let mut body = json!({ "model": effective_model, "max_tokens": cfg.max_output_tokens,
|
||
"system": system_prompt, "messages": messages });
|
||
if let Some(e) = effort {
|
||
let (thinking, output_config) =
|
||
crate::config::anthropic_thinking_config(effective_model, e, cfg.max_output_tokens);
|
||
if let Some(t) = thinking {
|
||
body["thinking"] = t;
|
||
}
|
||
if let Some(oc) = output_config {
|
||
body["output_config"] = oc;
|
||
}
|
||
}
|
||
if !tools_json.is_empty() {
|
||
body["tools"] = Value::Array(tools_json);
|
||
}
|
||
body
|
||
}
|
||
|
||
fn anthropic_tool_result_content(content: &[ToolResultContent]) -> Vec<Value> {
|
||
content
|
||
.iter()
|
||
.map(|c| match c {
|
||
ToolResultContent::Text(text) => json!({ "type": "text", "text": text }),
|
||
ToolResultContent::Image { data, mime_type } => json!({
|
||
"type": "image",
|
||
"source": { "type": "base64", "media_type": mime_type, "data": data },
|
||
}),
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
fn openai_body(
|
||
cfg: &Config,
|
||
system_prompt: &str,
|
||
history: &[HistoryItem],
|
||
tools: &[ToolDef],
|
||
effective_model: &str,
|
||
effort: Option<ThinkingEffort>,
|
||
) -> Value {
|
||
let mut messages: Vec<Value> = vec![json!({ "role": "system", "content": system_prompt })];
|
||
// Images returned from tool calls ride on a trailing `role:"user"`
|
||
// message because OpenAI Chat's `role:"tool"` content is text-only. We
|
||
// batch them across a run of adjacent ToolResult items so that all
|
||
// `role:"tool"` messages stay contiguous — splitting them with a user
|
||
// turn breaks OpenAI-Chat-compatible frontends that translate back to
|
||
// Anthropic `tool_result` (notably Databricks model serving), since
|
||
// Anthropic requires every `tool_use` in one assistant turn to be
|
||
// answered by a single immediately-following user message.
|
||
let mut pending_images: Vec<Value> = Vec::new();
|
||
let flush_images = |messages: &mut Vec<Value>, pending: &mut Vec<Value>| {
|
||
if !pending.is_empty() {
|
||
messages.push(json!({ "role": "user", "content": std::mem::take(pending) }));
|
||
}
|
||
};
|
||
for item in history {
|
||
match item {
|
||
HistoryItem::User(text) => {
|
||
flush_images(&mut messages, &mut pending_images);
|
||
messages.push(json!({ "role": "user", "content": text }));
|
||
}
|
||
HistoryItem::Assistant { text, tool_calls } => {
|
||
flush_images(&mut messages, &mut pending_images);
|
||
let mut msg = serde_json::Map::new();
|
||
msg.insert("role".into(), json!("assistant"));
|
||
msg.insert("content".into(), json!(text.as_str()));
|
||
if !tool_calls.is_empty() {
|
||
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()) } })
|
||
})
|
||
.collect();
|
||
msg.insert("tool_calls".into(), Value::Array(calls));
|
||
}
|
||
messages.push(Value::Object(msg));
|
||
}
|
||
HistoryItem::ToolResult(r) => {
|
||
messages.push(json!({
|
||
"role": "tool", "tool_call_id": r.provider_id,
|
||
"content": openai_tool_text_content(&r.content) }));
|
||
pending_images.extend(openai_image_user_content(&r.content));
|
||
}
|
||
}
|
||
}
|
||
flush_images(&mut messages, &mut pending_images);
|
||
let tools_json: Vec<Value> = tools
|
||
.iter()
|
||
.map(|t| {
|
||
json!({
|
||
"type": "function",
|
||
"function": { "name": t.name, "description": t.description,
|
||
"parameters": t.input_schema } })
|
||
})
|
||
.collect();
|
||
let mut body = json!({ "model": effective_model, "stream": false,
|
||
"max_completion_tokens": cfg.max_output_tokens, "messages": messages });
|
||
if let Some(e) = effort {
|
||
body["reasoning_effort"] = json!(e.openai_effort_str());
|
||
}
|
||
if !tools_json.is_empty() {
|
||
body["tools"] = Value::Array(tools_json);
|
||
body["tool_choice"] = json!("auto");
|
||
}
|
||
body
|
||
}
|
||
|
||
fn openai_tool_text_content(content: &[ToolResultContent]) -> String {
|
||
let mut parts = Vec::new();
|
||
for c in content {
|
||
match c {
|
||
ToolResultContent::Text(text) => parts.push(text.clone()),
|
||
ToolResultContent::Image { data, mime_type } => parts.push(format!(
|
||
"This tool result included an image ({mime_type}, {} base64 bytes) that is provided in the next user message.",
|
||
data.len()
|
||
)),
|
||
}
|
||
}
|
||
parts.join("\n")
|
||
}
|
||
|
||
fn openai_image_user_content(content: &[ToolResultContent]) -> Vec<Value> {
|
||
content
|
||
.iter()
|
||
.filter_map(|c| match c {
|
||
ToolResultContent::Image { data, mime_type } => Some(json!({
|
||
"type": "image_url",
|
||
"image_url": { "url": format!("data:{mime_type};base64,{data}") },
|
||
})),
|
||
ToolResultContent::Text(_) => None,
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
// Spec: https://platform.openai.com/docs/api-reference/responses
|
||
//
|
||
// Replay invariant: each assistant `function_call` input item **must**
|
||
// precede its matching `function_call_output`, or the API rejects with
|
||
// "No tool call found for call_id ...". `HistoryItem` ordering already
|
||
// guarantees this.
|
||
|
||
fn responses_body(
|
||
cfg: &Config,
|
||
system_prompt: &str,
|
||
history: &[HistoryItem],
|
||
tools: &[ToolDef],
|
||
effective_model: &str,
|
||
effort: Option<ThinkingEffort>,
|
||
) -> Value {
|
||
let mut input: Vec<Value> = Vec::with_capacity(history.len());
|
||
for item in history {
|
||
match item {
|
||
HistoryItem::User(text) => input.push(json!({
|
||
"role": "user",
|
||
"content": [{ "type": "input_text", "text": text }],
|
||
})),
|
||
HistoryItem::Assistant { text, tool_calls } => {
|
||
if !text.is_empty() {
|
||
input.push(json!({
|
||
"role": "assistant",
|
||
"content": [{ "type": "output_text", "text": text }],
|
||
}));
|
||
}
|
||
for c in tool_calls {
|
||
input.push(json!({
|
||
"type": "function_call",
|
||
"call_id": c.provider_id,
|
||
"name": c.name,
|
||
"arguments": serde_json::to_string(&c.arguments)
|
||
.unwrap_or_else(|_| "{}".into()),
|
||
}));
|
||
}
|
||
}
|
||
HistoryItem::ToolResult(r) => {
|
||
input.push(json!({
|
||
"type": "function_call_output",
|
||
"call_id": r.provider_id,
|
||
"output": openai_tool_text_content(&r.content),
|
||
}));
|
||
// Responses takes images as `input_image` parts on a user message.
|
||
let images: Vec<Value> = r
|
||
.content
|
||
.iter()
|
||
.filter_map(|c| match c {
|
||
ToolResultContent::Image { data, mime_type } => Some(json!({
|
||
"type": "input_image",
|
||
"image_url": format!("data:{mime_type};base64,{data}"),
|
||
})),
|
||
ToolResultContent::Text(_) => None,
|
||
})
|
||
.collect();
|
||
if !images.is_empty() {
|
||
input.push(json!({ "role": "user", "content": images }));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
let tools_json: Vec<Value> = tools
|
||
.iter()
|
||
.map(|t| {
|
||
json!({
|
||
"type": "function",
|
||
"name": t.name,
|
||
"description": t.description,
|
||
"parameters": t.input_schema,
|
||
})
|
||
})
|
||
.collect();
|
||
|
||
let mut body = json!({
|
||
"model": effective_model,
|
||
"instructions": system_prompt,
|
||
"max_output_tokens": cfg.max_output_tokens,
|
||
"input": input,
|
||
});
|
||
if let Some(e) = effort {
|
||
body["reasoning"] = json!({ "effort": e.openai_effort_str() });
|
||
}
|
||
if !tools_json.is_empty() {
|
||
body["tools"] = Value::Array(tools_json);
|
||
body["tool_choice"] = json!("auto");
|
||
}
|
||
body
|
||
}
|
||
|
||
/// Narrow matcher for "you should be on the Responses API" provider errors,
|
||
/// the signal we use to auto-upgrade. Triggers on the literal path
|
||
/// `/v1/responses` (Databricks GPT-5.5 phrasing) or the prose
|
||
/// "use the Responses API" / "Responses API instead".
|
||
fn is_responses_required_error(body: &str) -> bool {
|
||
let b = body.to_ascii_lowercase();
|
||
b.contains("/v1/responses")
|
||
|| b.contains("responses api instead")
|
||
|| b.contains("use the responses api")
|
||
}
|
||
|
||
fn databricks_v2_route_for_model(model: &str) -> DatabricksV2Route {
|
||
// Databricks v2 catalog names currently identify OpenAI-shaped GPT-5
|
||
// models and Anthropic-shaped Claude models by these substrings.
|
||
let lower = model.to_ascii_lowercase();
|
||
if lower.contains("gpt-5") || lower.contains("gpt5") {
|
||
DatabricksV2Route::OpenAiResponses
|
||
} else if lower.contains("claude") {
|
||
DatabricksV2Route::AnthropicMessages
|
||
} else {
|
||
DatabricksV2Route::MlflowChatCompletions
|
||
}
|
||
}
|
||
|
||
fn databricks_v2_path(route: DatabricksV2Route) -> &'static str {
|
||
match route {
|
||
DatabricksV2Route::OpenAiResponses => "/ai-gateway/openai/v1/responses",
|
||
DatabricksV2Route::AnthropicMessages => "/ai-gateway/anthropic/v1/messages",
|
||
DatabricksV2Route::MlflowChatCompletions => "/ai-gateway/mlflow/v1/chat/completions",
|
||
}
|
||
}
|
||
|
||
fn parse_responses(v: Value) -> Result<LlmResponse, AgentError> {
|
||
let mut text = String::new();
|
||
let mut reasoning = String::new();
|
||
let mut tool_calls = Vec::new();
|
||
let mut saw_function_call = false;
|
||
|
||
for item in v
|
||
.get("output")
|
||
.and_then(Value::as_array)
|
||
.into_iter()
|
||
.flatten()
|
||
{
|
||
match item.get("type").and_then(Value::as_str) {
|
||
Some("message") => {
|
||
for p in item
|
||
.get("content")
|
||
.and_then(Value::as_array)
|
||
.into_iter()
|
||
.flatten()
|
||
{
|
||
// Responses emits "output_text"; accept "text" forward-compat.
|
||
if matches!(
|
||
p.get("type").and_then(Value::as_str),
|
||
Some("output_text" | "text")
|
||
) {
|
||
if let Some(t) = p.get("text").and_then(Value::as_str) {
|
||
text.push_str(t);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
Some("function_call") => {
|
||
saw_function_call = true;
|
||
let raw = item
|
||
.get("arguments")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or("{}");
|
||
let args: Value = serde_json::from_str(raw).map_err(|e| {
|
||
AgentError::Llm(format!("function_call.arguments not valid JSON: {e}"))
|
||
})?;
|
||
tool_calls.push(make_tool_call(
|
||
str_field(item, "call_id"),
|
||
str_field(item, "name"),
|
||
args,
|
||
)?);
|
||
}
|
||
Some("reasoning") => {
|
||
// Reasoning summary items from the Responses API. Each item has a
|
||
// `summary` array of `{"type": "summary_text", "text": "..."}` objects.
|
||
for s in item
|
||
.get("summary")
|
||
.and_then(Value::as_array)
|
||
.into_iter()
|
||
.flatten()
|
||
{
|
||
if matches!(
|
||
s.get("type").and_then(Value::as_str),
|
||
Some("summary_text" | "text")
|
||
) {
|
||
if let Some(t) = s.get("text").and_then(Value::as_str) {
|
||
if !reasoning.is_empty() {
|
||
reasoning.push('\n');
|
||
}
|
||
reasoning.push_str(t);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
// Unknown types ignored for forward-compat.
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
let stop = match v.get("status").and_then(Value::as_str) {
|
||
Some("incomplete") => {
|
||
let reason = v
|
||
.get("incomplete_details")
|
||
.and_then(|d| d.get("reason"))
|
||
.and_then(Value::as_str);
|
||
if reason == Some("max_output_tokens") {
|
||
ProviderStop::MaxTokens
|
||
} else {
|
||
ProviderStop::Other
|
||
}
|
||
}
|
||
Some("completed") if saw_function_call => ProviderStop::ToolUse,
|
||
Some("completed") => ProviderStop::EndTurn,
|
||
_ => ProviderStop::Other,
|
||
};
|
||
let input_tokens = sum_usage(&v, &["input_tokens"]);
|
||
let output_tokens = sum_usage(&v, &["output_tokens"]);
|
||
Ok(LlmResponse {
|
||
text,
|
||
tool_calls,
|
||
stop,
|
||
input_tokens,
|
||
output_tokens,
|
||
reasoning,
|
||
})
|
||
}
|
||
|
||
fn map_stop(s: Option<&str>) -> ProviderStop {
|
||
match s {
|
||
Some("end_turn" | "stop") => ProviderStop::EndTurn,
|
||
Some("tool_use" | "tool_calls") => ProviderStop::ToolUse,
|
||
Some("max_tokens" | "length") => ProviderStop::MaxTokens,
|
||
Some("refusal" | "content_filter") => ProviderStop::Refusal,
|
||
_ => ProviderStop::Other,
|
||
}
|
||
}
|
||
|
||
/// Sum a set of `usage` token fields, returning `None` only when the `usage`
|
||
/// object is absent or carries none of the requested fields. A field that is
|
||
/// present is added; a field that is missing contributes 0. This keeps the
|
||
/// result an inclusive total (so cached tokens are never silently dropped)
|
||
/// while still distinguishing "no usage reported" from "usage was zero".
|
||
fn sum_usage(v: &Value, fields: &[&str]) -> Option<u64> {
|
||
let usage = v.get("usage")?;
|
||
let mut total: u64 = 0;
|
||
let mut saw_any = false;
|
||
for f in fields {
|
||
if let Some(n) = usage.get(*f).and_then(Value::as_u64) {
|
||
total = total.saturating_add(n);
|
||
saw_any = true;
|
||
}
|
||
}
|
||
saw_any.then_some(total)
|
||
}
|
||
|
||
/// Input-token total for Anthropic / Databricks (Anthropic-style) responses.
|
||
/// `input_tokens` alone EXCLUDES cached tokens, so we sum it with the two
|
||
/// cache fields to get the inclusive total the context budget must gate on.
|
||
fn anthropic_input_tokens(v: &Value) -> Option<u64> {
|
||
sum_usage(
|
||
v,
|
||
&[
|
||
"input_tokens",
|
||
"cache_read_input_tokens",
|
||
"cache_creation_input_tokens",
|
||
],
|
||
)
|
||
}
|
||
|
||
/// Input-token total for OpenAI Chat Completions and Databricks responses.
|
||
/// OpenAI's `prompt_tokens` is already inclusive. Databricks uses the same
|
||
/// `prompt_tokens` wire field but ALSO reports Anthropic-style cache fields
|
||
/// alongside it, so we sum them; the cache fields are simply absent (and
|
||
/// contribute 0) for vanilla OpenAI.
|
||
fn openai_chat_input_tokens(v: &Value) -> Option<u64> {
|
||
sum_usage(
|
||
v,
|
||
&[
|
||
"prompt_tokens",
|
||
"cache_read_input_tokens",
|
||
"cache_creation_input_tokens",
|
||
],
|
||
)
|
||
}
|
||
|
||
fn str_field(v: &Value, key: &str) -> String {
|
||
v.get(key).and_then(Value::as_str).unwrap_or("").to_owned()
|
||
}
|
||
|
||
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();
|
||
let mut text = String::new();
|
||
let mut reasoning = String::new();
|
||
if let Some(blocks) = v.get("content").and_then(Value::as_array) {
|
||
for b in blocks {
|
||
match b.get("type").and_then(Value::as_str) {
|
||
Some("text") => {
|
||
if let Some(t) = b.get("text").and_then(Value::as_str) {
|
||
text.push_str(t);
|
||
}
|
||
}
|
||
Some("thinking") => {
|
||
// Anthropic extended thinking block: `{"type": "thinking", "thinking": "..."}`
|
||
if let Some(t) = b.get("thinking").and_then(Value::as_str) {
|
||
if !reasoning.is_empty() {
|
||
reasoning.push('\n');
|
||
}
|
||
reasoning.push_str(t);
|
||
}
|
||
}
|
||
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),
|
||
)?),
|
||
_ => {}
|
||
}
|
||
}
|
||
}
|
||
let input_tokens = anthropic_input_tokens(&v);
|
||
let output_tokens = sum_usage(&v, &["output_tokens"]);
|
||
Ok(LlmResponse {
|
||
text,
|
||
tool_calls,
|
||
stop,
|
||
input_tokens,
|
||
output_tokens,
|
||
reasoning,
|
||
})
|
||
}
|
||
|
||
fn parse_openai(v: Value) -> Result<LlmResponse, AgentError> {
|
||
let choice = v
|
||
.get("choices")
|
||
.and_then(Value::as_array)
|
||
.and_then(|a| a.first())
|
||
.ok_or_else(|| AgentError::Llm("response missing choices".into()))?;
|
||
let stop = map_stop(choice.get("finish_reason").and_then(Value::as_str));
|
||
let msg = choice
|
||
.get("message")
|
||
.ok_or_else(|| AgentError::Llm("missing message".into()))?;
|
||
let text = str_field(msg, "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.
|
||
let reasoning = {
|
||
let rc = str_field(msg, "reasoning_content");
|
||
if rc.is_empty() {
|
||
str_field(msg, "reasoning")
|
||
} else {
|
||
rc
|
||
}
|
||
};
|
||
let mut tool_calls = Vec::new();
|
||
if let Some(arr) = msg.get("tool_calls").and_then(Value::as_array) {
|
||
for tc in arr {
|
||
let f = tc
|
||
.get("function")
|
||
.ok_or_else(|| AgentError::Llm("tool_call missing function".into()))?;
|
||
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}")))?;
|
||
tool_calls.push(make_tool_call(
|
||
str_field(tc, "id"),
|
||
str_field(f, "name"),
|
||
args,
|
||
)?);
|
||
}
|
||
}
|
||
let input_tokens = openai_chat_input_tokens(&v);
|
||
let output_tokens = sum_usage(&v, &["completion_tokens"]);
|
||
Ok(LlmResponse {
|
||
text,
|
||
tool_calls,
|
||
stop,
|
||
input_tokens,
|
||
output_tokens,
|
||
reasoning,
|
||
})
|
||
}
|
||
|
||
fn make_tool_call(id: String, name: String, args: Value) -> Result<ToolCall, AgentError> {
|
||
if id.is_empty() || name.is_empty() {
|
||
return Err(AgentError::Llm("tool_call missing id or name".into()));
|
||
}
|
||
let arguments = match args {
|
||
Value::Object(_) => args,
|
||
Value::Null => Value::Object(Default::default()),
|
||
_ => {
|
||
return Err(AgentError::Llm(
|
||
"tool_call arguments must be a JSON object".into(),
|
||
))
|
||
}
|
||
};
|
||
Ok(ToolCall {
|
||
provider_id: id,
|
||
name,
|
||
arguments,
|
||
})
|
||
}
|
||
|
||
async fn read_error_body(mut resp: reqwest::Response) -> String {
|
||
let mut buf: Vec<u8> = Vec::new();
|
||
while buf.len() < MAX_LLM_ERROR_BODY_BYTES {
|
||
match resp.chunk().await {
|
||
Ok(Some(chunk)) => {
|
||
let take = chunk.len().min(MAX_LLM_ERROR_BODY_BYTES - buf.len());
|
||
buf.extend_from_slice(&chunk[..take]);
|
||
if take < chunk.len() {
|
||
break;
|
||
}
|
||
}
|
||
_ => break,
|
||
}
|
||
}
|
||
String::from_utf8_lossy(&buf).into_owned()
|
||
}
|
||
|
||
const MAX_RETRIES: u32 = 3;
|
||
const BASE_BACKOFF_MS: u64 = 500;
|
||
const MAX_BACKOFF_MS: u64 = 8_000;
|
||
|
||
async fn backoff_with_jitter(attempt: u32) {
|
||
let base = BASE_BACKOFF_MS
|
||
.saturating_mul(1u64 << attempt)
|
||
.min(MAX_BACKOFF_MS);
|
||
let mut buf = [0u8; 8];
|
||
let jitter_range = base / 2;
|
||
let delay = if jitter_range > 0 && getrandom::fill(&mut buf).is_ok() {
|
||
let r = u64::from_le_bytes(buf) % jitter_range;
|
||
base - jitter_range + r
|
||
} else {
|
||
base
|
||
};
|
||
tokio::time::sleep(std::time::Duration::from_millis(delay)).await;
|
||
}
|
||
|
||
/// Transport-layer errors safe to retry for non-streaming LLM POSTs.
|
||
///
|
||
/// Covers timeouts, connect failures, and the broader request-class errors
|
||
/// reqwest reports for pre-response failures: TLS handshake aborts, sockets
|
||
/// dropped or reset mid-send, h2 GOAWAY/RST_STREAM, hyper protocol errors.
|
||
/// Body-serialization happens before the retry loop, so `is_request()` here
|
||
/// is always a network failure, never a malformed request we'd just resend.
|
||
fn is_retryable_transport_error(e: &reqwest::Error) -> bool {
|
||
e.is_timeout() || e.is_connect() || e.is_request()
|
||
}
|
||
|
||
/// Build the terminal `AgentError::Llm` for a `post()` exit that has given up
|
||
/// retrying — persistent retryable status, transport failure, or a body-read
|
||
/// break. `detail` carries the specific cause (status/body, or the transport
|
||
/// error text); `elapsed` and `attempts` are the cumulative cost of every
|
||
/// attempt made on this call, including the retries that failed before this
|
||
/// one. When cumulative time crosses `STALL_NOTICE_THRESHOLD`, this also logs
|
||
/// a `tracing::warn!` so a slow-building outage is visible in logs even
|
||
/// before an operator reads the returned error text.
|
||
fn terminal_llm_error(elapsed: std::time::Duration, attempts: u32, detail: &str) -> AgentError {
|
||
if elapsed >= STALL_NOTICE_THRESHOLD {
|
||
tracing::warn!(
|
||
cumulative_stall = ?elapsed,
|
||
attempts,
|
||
"llm: cumulative stall {elapsed:?} across {attempts} attempts ({detail})"
|
||
);
|
||
}
|
||
AgentError::Llm(format!(
|
||
"{detail} (cumulative {elapsed:?}, {attempts} attempt{})",
|
||
if attempts == 1 { "" } else { "s" },
|
||
))
|
||
}
|
||
|
||
/// Internal HTTP failure that preserves a mesh-specific MoA failure as a
|
||
/// typed signal. This covers both pre-inference eligibility rejection and a
|
||
/// structured `moa_failure` from workers/reducers. It is consumed inside
|
||
/// `openai_request`: an adaptive `auto` call falls back once, while an explicit
|
||
/// `mesh` call is surfaced as the ordinary LLM error callers already
|
||
/// understand.
|
||
#[derive(Debug)]
|
||
enum PostError {
|
||
Agent(AgentError),
|
||
MeshFallback(String),
|
||
}
|
||
|
||
impl PostError {
|
||
fn into_agent(self) -> AgentError {
|
||
match self {
|
||
Self::Agent(error) => error,
|
||
Self::MeshFallback(detail) => AgentError::Llm(detail),
|
||
}
|
||
}
|
||
}
|
||
|
||
impl From<AgentError> for PostError {
|
||
fn from(error: AgentError) -> Self {
|
||
Self::Agent(error)
|
||
}
|
||
}
|
||
|
||
fn is_mesh_moa_unavailable_body(body: &str) -> bool {
|
||
serde_json::from_str::<Value>(body)
|
||
.ok()
|
||
.and_then(|value| {
|
||
value
|
||
.pointer("/error/message")
|
||
.and_then(Value::as_str)
|
||
.map(str::to_owned)
|
||
})
|
||
.as_deref()
|
||
== Some(MESH_MOA_UNAVAILABLE_MESSAGE)
|
||
}
|
||
|
||
fn is_mesh_moa_failure_body(body: &str) -> bool {
|
||
serde_json::from_str::<Value>(body)
|
||
.ok()
|
||
.and_then(|value| {
|
||
value
|
||
.pointer("/error/type")
|
||
.and_then(Value::as_str)
|
||
.map(str::to_owned)
|
||
})
|
||
.as_deref()
|
||
== Some("moa_failure")
|
||
}
|
||
|
||
async fn post<F>(
|
||
http: &Client,
|
||
url: &str,
|
||
body: &Value,
|
||
detect_mesh_fallback: bool,
|
||
apply: F,
|
||
) -> Result<Value, PostError>
|
||
where
|
||
F: Fn(reqwest::RequestBuilder) -> reqwest::RequestBuilder,
|
||
{
|
||
let body_bytes = serde_json::to_vec(body)
|
||
.map_err(|e| PostError::Agent(AgentError::Llm(format!("serialize: {e}"))))?;
|
||
let call_start = std::time::Instant::now();
|
||
for attempt in 0..MAX_RETRIES {
|
||
let resp = match apply(
|
||
http.post(url)
|
||
.header("content-type", "application/json")
|
||
.body(body_bytes.clone()),
|
||
)
|
||
.send()
|
||
.await
|
||
{
|
||
Ok(r) => r,
|
||
Err(e) => {
|
||
if attempt + 1 < MAX_RETRIES && is_retryable_transport_error(&e) {
|
||
tracing::warn!(
|
||
attempt = attempt + 1,
|
||
max_attempts = MAX_RETRIES,
|
||
error = %e,
|
||
"llm: transport error, retrying"
|
||
);
|
||
backoff_with_jitter(attempt).await;
|
||
continue;
|
||
}
|
||
return Err(PostError::Agent(terminal_llm_error(
|
||
call_start.elapsed(),
|
||
attempt + 1,
|
||
&format!("transport: {e}"),
|
||
)));
|
||
}
|
||
};
|
||
let status = resp.status();
|
||
// Both 401 and 403 are treated as refreshable: a 403 can mean an
|
||
// expired or revoked token, not just a pure authorization verdict, and
|
||
// the two are indistinguishable at the HTTP-status layer. The caller's
|
||
// retry loop keys off `LlmAuth` and refreshes once; the per-call guard
|
||
// bounds a pure-authz 403 to one wasted refresh before it propagates.
|
||
// Not a stall path: auth failures are not surfaced through
|
||
// terminal_llm_error, they resolve on the next call after refresh.
|
||
if status == 401 || status == 403 {
|
||
return Err(PostError::Agent(AgentError::LlmAuth(
|
||
read_error_body(resp).await,
|
||
)));
|
||
}
|
||
if status.is_server_error() || status == 429 || status.as_u16() == 499 {
|
||
let body = read_error_body(resp).await;
|
||
let mesh_fallback = detect_mesh_fallback
|
||
&& (status == reqwest::StatusCode::SERVICE_UNAVAILABLE
|
||
&& is_mesh_moa_unavailable_body(&body)
|
||
|| status.is_server_error() && is_mesh_moa_failure_body(&body));
|
||
if mesh_fallback {
|
||
return Err(PostError::MeshFallback(format!("{status}: {body}")));
|
||
}
|
||
if attempt + 1 < MAX_RETRIES {
|
||
tracing::warn!(
|
||
attempt = attempt + 1,
|
||
max_attempts = MAX_RETRIES,
|
||
%status,
|
||
"llm: retryable status, retrying"
|
||
);
|
||
backoff_with_jitter(attempt).await;
|
||
continue;
|
||
}
|
||
return Err(PostError::Agent(terminal_llm_error(
|
||
call_start.elapsed(),
|
||
attempt + 1,
|
||
&format!("exhausted retries: {status}: {body}"),
|
||
)));
|
||
}
|
||
// Not a stall path: the model is misconfigured, not the transport or
|
||
// upstream capacity — no retry was attempted, so cumulative duration
|
||
// would be misleading.
|
||
if status == 404 {
|
||
return Err(PostError::Agent(AgentError::LlmModelNotFound(format!(
|
||
"{status}: {}",
|
||
read_error_body(resp).await
|
||
))));
|
||
}
|
||
if !status.is_success() {
|
||
return Err(PostError::Agent(AgentError::Llm(format!(
|
||
"{status}: {}",
|
||
read_error_body(resp).await
|
||
))));
|
||
}
|
||
if let Some(len) = resp.content_length() {
|
||
if len as usize > MAX_LLM_RESPONSE_BYTES {
|
||
return Err(PostError::Agent(AgentError::Llm(format!(
|
||
"response too large: {len} > {MAX_LLM_RESPONSE_BYTES}"
|
||
))));
|
||
}
|
||
}
|
||
let mut buf: Vec<u8> = Vec::new();
|
||
let mut stream = resp;
|
||
loop {
|
||
match stream.chunk().await {
|
||
Ok(Some(chunk)) => {
|
||
if buf.len() + chunk.len() > MAX_LLM_RESPONSE_BYTES {
|
||
return Err(PostError::Agent(AgentError::Llm(format!(
|
||
"response exceeded {MAX_LLM_RESPONSE_BYTES} bytes"
|
||
))));
|
||
}
|
||
buf.extend_from_slice(&chunk);
|
||
}
|
||
Ok(None) => break,
|
||
Err(e) => {
|
||
return Err(PostError::Agent(terminal_llm_error(
|
||
call_start.elapsed(),
|
||
attempt + 1,
|
||
&format!("body read: {e}"),
|
||
)));
|
||
}
|
||
}
|
||
}
|
||
return serde_json::from_slice(&buf)
|
||
.map_err(|e| PostError::Agent(AgentError::Llm(format!("json: {e}"))));
|
||
}
|
||
unreachable!("loop always returns on its final iteration (attempt + 1 == MAX_RETRIES)");
|
||
}
|
||
|
||
/// Build the `TokenSource` for the configured provider.
|
||
///
|
||
/// - `Provider::Anthropic`: a static source seeded from `cfg.api_key`. It's
|
||
/// never read for Anthropic requests (those go through `post_anthropic` with
|
||
/// `x-api-key`), but Llm holds one to keep the field non-`Option`.
|
||
/// - `Provider::OpenAi`: a static source over `OPENAI_COMPAT_API_KEY`.
|
||
/// - `Provider::Databricks`: if `DATABRICKS_TOKEN` is set, a static source.
|
||
/// Otherwise a `PkceOAuthTokenSource` pointed at the workspace's OIDC
|
||
/// discovery URL. First request without a cached token triggers a browser
|
||
/// flow; subsequent requests use the cache + refresh transparently.
|
||
pub(crate) fn build_token_source(cfg: &Config) -> Result<Arc<dyn TokenSource>, AgentError> {
|
||
match cfg.provider {
|
||
Provider::Anthropic | Provider::OpenAi => {
|
||
Ok(Arc::new(StaticTokenSource::new(cfg.api_key.clone())))
|
||
}
|
||
Provider::Databricks | Provider::DatabricksV2 => {
|
||
if !cfg.api_key.is_empty() {
|
||
return Ok(Arc::new(StaticTokenSource::new(cfg.api_key.clone())));
|
||
}
|
||
let discovery_url = format!(
|
||
"{}/oidc/.well-known/oauth-authorization-server",
|
||
cfg.base_url.trim_end_matches('/')
|
||
);
|
||
let pkce = PkceOAuthConfig {
|
||
discovery_url,
|
||
client_id: DATABRICKS_CLIENT_ID.into(),
|
||
scopes: DATABRICKS_OAUTH_SCOPES
|
||
.iter()
|
||
.map(|s| (*s).into())
|
||
.collect(),
|
||
cache_namespace: "databricks".into(),
|
||
cache_dir_override: None,
|
||
};
|
||
Ok(PkceOAuthTokenSource::new(pkce)?)
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Return a clone of `body` with any top-level `"model"` field removed.
|
||
/// Used for Databricks model-serving, which encodes the model in the URL
|
||
/// path and rejects the field in the body.
|
||
fn strip_model(body: &Value) -> Value {
|
||
match body {
|
||
Value::Object(map) => {
|
||
let mut m = map.clone();
|
||
m.remove("model");
|
||
Value::Object(m)
|
||
}
|
||
other => other.clone(),
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::config::{Config, HookServers, OpenAiApi, Provider};
|
||
use crate::types::{HistoryItem, ToolCall, ToolResult, ToolResultContent};
|
||
use std::time::Duration;
|
||
use tracing_subscriber::layer::SubscriberExt;
|
||
|
||
fn cfg(provider: Provider) -> Config {
|
||
Config {
|
||
provider,
|
||
system_prompt: "system".into(),
|
||
max_rounds: 10,
|
||
max_output_tokens: 1024,
|
||
llm_timeout: Duration::from_secs(10),
|
||
tool_timeout: Duration::from_secs(10),
|
||
mcp_init_timeout: Duration::from_secs(10),
|
||
mcp_max_restart_attempts: 1,
|
||
mcp_restart_base_ms: 1,
|
||
mcp_restart_max_ms: 1,
|
||
max_sessions: 1,
|
||
max_line_bytes: 1024 * 1024,
|
||
max_history_bytes: 16 * 1024 * 1024,
|
||
max_tool_result_text_bytes: 50 * 1024,
|
||
max_context_tokens: 200_000,
|
||
max_handoffs: 1,
|
||
max_parallel_tools: 1,
|
||
hook_timeout: Duration::from_secs(1),
|
||
stop_max_rejections: 0,
|
||
hook_servers: HookServers::None,
|
||
api_key: "key".into(),
|
||
model: "model".into(),
|
||
base_url: "http://example.invalid".into(),
|
||
anthropic_api_version: "2023-06-01".into(),
|
||
openai_api: OpenAiApi::Chat,
|
||
prefer_mesh_for_auto: false,
|
||
hints_enabled: true,
|
||
thinking_effort: None,
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
struct CapturedHttpRequest {
|
||
method: String,
|
||
path: String,
|
||
body: Option<Value>,
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
struct StubHttpResponse {
|
||
status: u16,
|
||
body: Value,
|
||
}
|
||
|
||
impl StubHttpResponse {
|
||
fn ok(body: Value) -> Self {
|
||
Self { status: 200, body }
|
||
}
|
||
|
||
fn error(status: u16, message: &str) -> Self {
|
||
Self {
|
||
status,
|
||
body: json!({
|
||
"error": {
|
||
"message": message,
|
||
"type": "server_error",
|
||
"code": "service_unavailable",
|
||
}
|
||
}),
|
||
}
|
||
}
|
||
|
||
fn moa_failure(status: u16, code: &str, message: &str) -> Self {
|
||
Self {
|
||
status,
|
||
body: json!({
|
||
"choices": [{
|
||
"finish_reason": "error",
|
||
"message": { "content": message, "role": "assistant" },
|
||
}],
|
||
"error": {
|
||
"message": message,
|
||
"type": "moa_failure",
|
||
"code": code,
|
||
},
|
||
"model": "mesh",
|
||
}),
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn spawn_sequence_stub(
|
||
responses: Vec<StubHttpResponse>,
|
||
) -> (String, Arc<Mutex<Vec<CapturedHttpRequest>>>) {
|
||
use std::collections::VecDeque;
|
||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||
use tokio::net::TcpListener;
|
||
|
||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||
let base_url = format!("http://{}/v1", listener.local_addr().unwrap());
|
||
let captured = Arc::new(Mutex::new(Vec::new()));
|
||
let captured_for_server = captured.clone();
|
||
let responses = Arc::new(Mutex::new(VecDeque::from(responses)));
|
||
|
||
tokio::spawn(async move {
|
||
loop {
|
||
let (mut socket, _) = match listener.accept().await {
|
||
Ok(connection) => connection,
|
||
Err(_) => return,
|
||
};
|
||
let mut bytes = Vec::new();
|
||
let mut chunk = [0u8; 4096];
|
||
let header_end = loop {
|
||
if let Some(index) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
|
||
break index + 4;
|
||
}
|
||
match socket.read(&mut chunk).await {
|
||
Ok(0) | Err(_) => return,
|
||
Ok(read) => bytes.extend_from_slice(&chunk[..read]),
|
||
}
|
||
};
|
||
let header_text = String::from_utf8_lossy(&bytes[..header_end]).into_owned();
|
||
let content_length = header_text
|
||
.lines()
|
||
.find_map(|line| {
|
||
let (name, value) = line.split_once(':')?;
|
||
name.eq_ignore_ascii_case("content-length")
|
||
.then(|| value.trim().parse::<usize>().ok())
|
||
.flatten()
|
||
})
|
||
.unwrap_or(0);
|
||
while bytes.len() < header_end + content_length {
|
||
match socket.read(&mut chunk).await {
|
||
Ok(0) | Err(_) => return,
|
||
Ok(read) => bytes.extend_from_slice(&chunk[..read]),
|
||
}
|
||
}
|
||
let mut request_line = header_text.lines().next().unwrap_or_default().split(' ');
|
||
let method = request_line.next().unwrap_or_default().to_string();
|
||
let path = request_line.next().unwrap_or_default().to_string();
|
||
let body = if content_length == 0 {
|
||
None
|
||
} else {
|
||
serde_json::from_slice(&bytes[header_end..header_end + content_length]).ok()
|
||
};
|
||
captured_for_server
|
||
.lock()
|
||
.await
|
||
.push(CapturedHttpRequest { method, path, body });
|
||
|
||
let response = responses.lock().await.pop_front().unwrap_or_else(|| {
|
||
StubHttpResponse::error(500, "stub response sequence exhausted")
|
||
});
|
||
let status_text = match response.status {
|
||
200 => "OK",
|
||
500 => "Internal Server Error",
|
||
502 => "Bad Gateway",
|
||
503 => "Service Unavailable",
|
||
other => panic!("unsupported stub status: {other}"),
|
||
};
|
||
let body = response.body.to_string();
|
||
let wire = format!(
|
||
"HTTP/1.1 {} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||
response.status,
|
||
status_text,
|
||
body.len(),
|
||
body,
|
||
);
|
||
let _ = socket.write_all(wire.as_bytes()).await;
|
||
let _ = socket.shutdown().await;
|
||
}
|
||
});
|
||
(base_url, captured)
|
||
}
|
||
|
||
fn model_catalog(ids: &[&str]) -> Value {
|
||
json!({
|
||
"object": "list",
|
||
"data": ids
|
||
.iter()
|
||
.map(|id| json!({ "id": id, "object": "model" }))
|
||
.collect::<Vec<_>>(),
|
||
})
|
||
}
|
||
|
||
fn chat_response(text: &str) -> Value {
|
||
json!({
|
||
"choices": [{
|
||
"finish_reason": "stop",
|
||
"message": { "content": text },
|
||
}]
|
||
})
|
||
}
|
||
|
||
async fn complete_model(
|
||
llm: &Llm,
|
||
cfg: &Config,
|
||
model: &str,
|
||
) -> Result<LlmResponse, AgentError> {
|
||
llm.complete(
|
||
cfg,
|
||
"system",
|
||
&[HistoryItem::User("hello".into())],
|
||
&[],
|
||
model,
|
||
)
|
||
.await
|
||
}
|
||
|
||
async fn complete_model_with_tool(
|
||
llm: &Llm,
|
||
cfg: &Config,
|
||
model: &str,
|
||
) -> Result<LlmResponse, AgentError> {
|
||
llm.complete(
|
||
cfg,
|
||
"system",
|
||
&[HistoryItem::User("add two numbers".into())],
|
||
&[ToolDef {
|
||
name: "add_numbers".into(),
|
||
description: "Add two numbers".into(),
|
||
input_schema: json!({
|
||
"type": "object",
|
||
"properties": {
|
||
"a": {"type": "number"},
|
||
"b": {"type": "number"}
|
||
},
|
||
"required": ["a", "b"]
|
||
}),
|
||
}],
|
||
model,
|
||
)
|
||
.await
|
||
}
|
||
|
||
async fn expire_mesh_catalog_check(llm: &Llm) {
|
||
llm.mesh_auto_state.lock().await.last_checked =
|
||
Some(Instant::now() - MESH_AUTO_CATALOG_TTL);
|
||
}
|
||
|
||
async fn expire_mesh_auto_cooldown(llm: &Llm) {
|
||
let mut state = llm.mesh_auto_state.lock().await;
|
||
state.last_checked = Some(Instant::now() - MESH_AUTO_CATALOG_TTL);
|
||
state.cooldown_until = Some(Instant::now() - std::time::Duration::from_secs(1));
|
||
}
|
||
|
||
fn posted_models(requests: &[CapturedHttpRequest]) -> Vec<&str> {
|
||
requests
|
||
.iter()
|
||
.filter(|request| request.method == "POST")
|
||
.filter_map(|request| request.body.as_ref()?.get("model")?.as_str())
|
||
.collect()
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn mesh_auto_requires_two_stable_catalog_observations() {
|
||
let (base_url, captured) = spawn_sequence_stub(vec![
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("direct")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("collective")),
|
||
])
|
||
.await;
|
||
let mut config = cfg(Provider::OpenAi);
|
||
config.base_url = base_url;
|
||
config.prefer_mesh_for_auto = true;
|
||
let llm = Llm::new(&config).unwrap();
|
||
|
||
assert_eq!(
|
||
complete_model(&llm, &config, "auto").await.unwrap().text,
|
||
"direct"
|
||
);
|
||
expire_mesh_catalog_check(&llm).await;
|
||
assert_eq!(
|
||
complete_model(&llm, &config, "auto").await.unwrap().text,
|
||
"collective"
|
||
);
|
||
|
||
let requests = captured.lock().await;
|
||
assert_eq!(posted_models(&requests), vec!["auto", "mesh"]);
|
||
assert_eq!(
|
||
requests
|
||
.iter()
|
||
.filter(|request| request.path == "/v1/models")
|
||
.count(),
|
||
2
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn mesh_auto_does_not_enable_while_second_model_flaps() {
|
||
let (base_url, captured) = spawn_sequence_stub(vec![
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("one")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a"])),
|
||
StubHttpResponse::ok(chat_response("two")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("three")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("four")),
|
||
])
|
||
.await;
|
||
let mut config = cfg(Provider::OpenAi);
|
||
config.base_url = base_url;
|
||
config.prefer_mesh_for_auto = true;
|
||
let llm = Llm::new(&config).unwrap();
|
||
|
||
for _ in 0..4 {
|
||
complete_model(&llm, &config, "auto").await.unwrap();
|
||
expire_mesh_catalog_check(&llm).await;
|
||
}
|
||
|
||
let requests = captured.lock().await;
|
||
assert_eq!(
|
||
posted_models(&requests),
|
||
vec!["auto", "auto", "auto", "mesh"]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn collective_catalog_requires_two_distinct_physical_models() {
|
||
assert_eq!(mesh_catalog_supports_collective(&json!({})), None);
|
||
assert_eq!(
|
||
mesh_catalog_supports_collective(&model_catalog(&[])),
|
||
Some(false)
|
||
);
|
||
assert_eq!(
|
||
mesh_catalog_supports_collective(&model_catalog(&["model-a", "mesh"])),
|
||
Some(false)
|
||
);
|
||
assert_eq!(
|
||
mesh_catalog_supports_collective(&model_catalog(&[
|
||
"org/model@main:Q4",
|
||
"org/model:Q4",
|
||
"mesh"
|
||
])),
|
||
Some(false),
|
||
"two spellings of one model are not collective capacity"
|
||
);
|
||
assert_eq!(
|
||
mesh_catalog_supports_collective(&model_catalog(&["model-a", "model-b", "mesh"])),
|
||
Some(true)
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn mesh_auto_tracks_models_appearing_disappearing_and_reappearing() {
|
||
let (base_url, captured) = spawn_sequence_stub(vec![
|
||
StubHttpResponse::ok(model_catalog(&[])),
|
||
StubHttpResponse::ok(chat_response("zero")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("one")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("two-first-observation")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("two-stable")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("contracted")),
|
||
StubHttpResponse::ok(model_catalog(&[])),
|
||
StubHttpResponse::ok(chat_response("empty-again")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("rejoin-first-observation")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("rejoined-stable")),
|
||
])
|
||
.await;
|
||
let mut config = cfg(Provider::OpenAi);
|
||
config.base_url = base_url;
|
||
config.prefer_mesh_for_auto = true;
|
||
let llm = Llm::new(&config).unwrap();
|
||
|
||
for call in 0..8 {
|
||
if call == 5 {
|
||
expire_mesh_auto_cooldown(&llm).await;
|
||
} else if call > 0 {
|
||
expire_mesh_catalog_check(&llm).await;
|
||
}
|
||
complete_model(&llm, &config, "auto").await.unwrap();
|
||
}
|
||
|
||
let requests = captured.lock().await;
|
||
assert_eq!(
|
||
posted_models(&requests),
|
||
vec!["auto", "auto", "auto", "mesh", "auto", "auto", "auto", "mesh"]
|
||
);
|
||
assert_eq!(
|
||
requests
|
||
.iter()
|
||
.filter(|request| request.path == "/v1/models")
|
||
.count(),
|
||
8
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn mesh_auto_falls_back_once_and_cools_down_when_mesh_contracts() {
|
||
let (base_url, captured) = spawn_sequence_stub(vec![
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("warmup")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::error(503, MESH_MOA_UNAVAILABLE_MESSAGE),
|
||
StubHttpResponse::ok(chat_response("fallback")),
|
||
StubHttpResponse::ok(chat_response("cooldown")),
|
||
])
|
||
.await;
|
||
let mut config = cfg(Provider::OpenAi);
|
||
config.base_url = base_url;
|
||
config.prefer_mesh_for_auto = true;
|
||
let llm = Llm::new(&config).unwrap();
|
||
|
||
complete_model(&llm, &config, "auto").await.unwrap();
|
||
expire_mesh_catalog_check(&llm).await;
|
||
assert_eq!(
|
||
complete_model(&llm, &config, "auto").await.unwrap().text,
|
||
"fallback"
|
||
);
|
||
assert_eq!(
|
||
complete_model(&llm, &config, "auto").await.unwrap().text,
|
||
"cooldown"
|
||
);
|
||
|
||
let requests = captured.lock().await;
|
||
assert_eq!(
|
||
posted_models(&requests),
|
||
vec!["auto", "mesh", "auto", "auto"]
|
||
);
|
||
assert_eq!(
|
||
requests
|
||
.iter()
|
||
.filter(|request| request.path == "/v1/models")
|
||
.count(),
|
||
2,
|
||
"cooldown request must not re-probe the catalog"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn mesh_auto_falls_back_once_when_moa_reducer_fails() {
|
||
let (base_url, captured) = spawn_sequence_stub(vec![
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("warmup")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::moa_failure(
|
||
502,
|
||
"all_reducers_failed",
|
||
"Reducer failed: unsupported chat template",
|
||
),
|
||
StubHttpResponse::ok(chat_response("fallback")),
|
||
StubHttpResponse::ok(chat_response("cooldown")),
|
||
])
|
||
.await;
|
||
let mut config = cfg(Provider::OpenAi);
|
||
config.base_url = base_url;
|
||
config.prefer_mesh_for_auto = true;
|
||
let llm = Llm::new(&config).unwrap();
|
||
|
||
complete_model(&llm, &config, "auto").await.unwrap();
|
||
expire_mesh_catalog_check(&llm).await;
|
||
assert_eq!(
|
||
complete_model(&llm, &config, "auto").await.unwrap().text,
|
||
"fallback"
|
||
);
|
||
assert_eq!(
|
||
complete_model(&llm, &config, "auto").await.unwrap().text,
|
||
"cooldown"
|
||
);
|
||
|
||
let requests = captured.lock().await;
|
||
assert_eq!(
|
||
posted_models(&requests),
|
||
vec!["auto", "mesh", "auto", "auto"]
|
||
);
|
||
assert_eq!(
|
||
requests
|
||
.iter()
|
||
.filter(|request| request.path == "/v1/models")
|
||
.count(),
|
||
2,
|
||
"mesh-specific failure must enter cooldown without re-probing"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn mesh_auto_retries_unstructured_tool_markup_through_auto() {
|
||
let (base_url, captured) = spawn_sequence_stub(vec![
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response("warmup")),
|
||
StubHttpResponse::ok(model_catalog(&["model-a", "model-b", "mesh"])),
|
||
StubHttpResponse::ok(chat_response(
|
||
"<|tool_call>call:add_numbers{a:17,b:25}<tool_call|>",
|
||
)),
|
||
StubHttpResponse::ok(chat_response("safe fallback")),
|
||
StubHttpResponse::ok(chat_response("cooldown")),
|
||
])
|
||
.await;
|
||
let mut config = cfg(Provider::OpenAi);
|
||
config.base_url = base_url;
|
||
config.prefer_mesh_for_auto = true;
|
||
let llm = Llm::new(&config).unwrap();
|
||
|
||
complete_model(&llm, &config, "auto").await.unwrap();
|
||
expire_mesh_catalog_check(&llm).await;
|
||
assert_eq!(
|
||
complete_model_with_tool(&llm, &config, "auto")
|
||
.await
|
||
.unwrap()
|
||
.text,
|
||
"safe fallback"
|
||
);
|
||
assert_eq!(
|
||
complete_model_with_tool(&llm, &config, "auto")
|
||
.await
|
||
.unwrap()
|
||
.text,
|
||
"cooldown"
|
||
);
|
||
|
||
let requests = captured.lock().await;
|
||
assert_eq!(
|
||
posted_models(&requests),
|
||
vec!["auto", "mesh", "auto", "auto"]
|
||
);
|
||
assert_eq!(
|
||
requests
|
||
.iter()
|
||
.filter(|request| request.path == "/v1/models")
|
||
.count(),
|
||
2,
|
||
"pseudo tool markup must enter cooldown without another catalog probe"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn mesh_auto_catalog_failure_fails_open_to_auto() {
|
||
let (base_url, captured) = spawn_sequence_stub(vec![
|
||
StubHttpResponse::error(500, "catalog unavailable"),
|
||
StubHttpResponse::ok(chat_response("direct")),
|
||
])
|
||
.await;
|
||
let mut config = cfg(Provider::OpenAi);
|
||
config.base_url = base_url;
|
||
config.prefer_mesh_for_auto = true;
|
||
let llm = Llm::new(&config).unwrap();
|
||
|
||
complete_model(&llm, &config, "auto").await.unwrap();
|
||
let requests = captured.lock().await;
|
||
assert_eq!(posted_models(&requests), vec!["auto"]);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn generic_openai_auto_does_not_probe_or_select_mesh() {
|
||
let (base_url, captured) =
|
||
spawn_sequence_stub(vec![StubHttpResponse::ok(chat_response("provider auto"))]).await;
|
||
let mut config = cfg(Provider::OpenAi);
|
||
config.base_url = base_url;
|
||
config.prefer_mesh_for_auto = false;
|
||
let llm = Llm::new(&config).unwrap();
|
||
|
||
complete_model(&llm, &config, "auto").await.unwrap();
|
||
let requests = captured.lock().await;
|
||
assert_eq!(posted_models(&requests), vec!["auto"]);
|
||
assert!(requests.iter().all(|request| request.path != "/v1/models"));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn explicit_models_are_never_rewritten_or_fallback_retried() {
|
||
let (base_url, captured) = spawn_sequence_stub(vec![StubHttpResponse::error(
|
||
503,
|
||
MESH_MOA_UNAVAILABLE_MESSAGE,
|
||
)])
|
||
.await;
|
||
let mut config = cfg(Provider::OpenAi);
|
||
config.base_url = base_url;
|
||
config.prefer_mesh_for_auto = true;
|
||
let llm = Llm::new(&config).unwrap();
|
||
|
||
let error = complete_model(&llm, &config, "mesh").await.unwrap_err();
|
||
assert!(error.to_string().contains(MESH_MOA_UNAVAILABLE_MESSAGE));
|
||
let requests = captured.lock().await;
|
||
assert_eq!(posted_models(&requests), vec!["mesh"]);
|
||
assert!(requests.iter().all(|request| request.path != "/v1/models"));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn explicit_real_model_bypasses_mesh_auto_policy() {
|
||
let (base_url, captured) =
|
||
spawn_sequence_stub(vec![StubHttpResponse::ok(chat_response("explicit"))]).await;
|
||
let mut config = cfg(Provider::OpenAi);
|
||
config.base_url = base_url;
|
||
config.prefer_mesh_for_auto = true;
|
||
let llm = Llm::new(&config).unwrap();
|
||
|
||
let response = complete_model(&llm, &config, "model-a").await.unwrap();
|
||
assert_eq!(response.text, "explicit");
|
||
let requests = captured.lock().await;
|
||
assert_eq!(posted_models(&requests), vec!["model-a"]);
|
||
assert!(requests.iter().all(|request| request.path != "/v1/models"));
|
||
}
|
||
|
||
#[test]
|
||
fn mesh_unavailable_classifier_requires_exact_openai_error_message() {
|
||
assert!(is_mesh_moa_unavailable_body(
|
||
&StubHttpResponse::error(503, MESH_MOA_UNAVAILABLE_MESSAGE)
|
||
.body
|
||
.to_string()
|
||
));
|
||
assert!(!is_mesh_moa_unavailable_body(
|
||
&StubHttpResponse::error(503, "some other outage")
|
||
.body
|
||
.to_string()
|
||
));
|
||
}
|
||
|
||
#[test]
|
||
fn mesh_failure_classifier_requires_structured_moa_failure_type() {
|
||
assert!(is_mesh_moa_failure_body(
|
||
&StubHttpResponse::moa_failure(
|
||
502,
|
||
"all_reducers_failed",
|
||
"Reducer failed: unsupported chat template",
|
||
)
|
||
.body
|
||
.to_string()
|
||
));
|
||
assert!(!is_mesh_moa_failure_body(
|
||
&StubHttpResponse::error(502, "some other bad gateway")
|
||
.body
|
||
.to_string()
|
||
));
|
||
}
|
||
|
||
fn image_history() -> Vec<HistoryItem> {
|
||
vec![
|
||
HistoryItem::User("describe the image".into()),
|
||
HistoryItem::Assistant {
|
||
text: String::new(),
|
||
tool_calls: vec![ToolCall {
|
||
provider_id: "toolu_1".into(),
|
||
name: "dev__view_image".into(),
|
||
arguments: serde_json::json!({"source":"x.png"}),
|
||
}],
|
||
},
|
||
HistoryItem::ToolResult(ToolResult {
|
||
provider_id: "toolu_1".into(),
|
||
content: vec![
|
||
ToolResultContent::Text("10×10, 70 B (image/png from x.png)".into()),
|
||
ToolResultContent::Image {
|
||
data: "aW1n".into(),
|
||
mime_type: "image/png".into(),
|
||
},
|
||
],
|
||
is_error: false,
|
||
}),
|
||
]
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_tool_result_preserves_image_block() {
|
||
let body = anthropic_body(
|
||
&cfg(Provider::Anthropic),
|
||
"system",
|
||
&image_history(),
|
||
&[],
|
||
"model",
|
||
None,
|
||
);
|
||
let content = &body["messages"][2]["content"][0]["content"];
|
||
assert_eq!(content[0]["type"], "text");
|
||
assert_eq!(content[1]["type"], "image");
|
||
assert_eq!(content[1]["source"]["type"], "base64");
|
||
assert_eq!(content[1]["source"]["media_type"], "image/png");
|
||
assert_eq!(content[1]["source"]["data"], "aW1n");
|
||
}
|
||
|
||
fn cfg_responses() -> Config {
|
||
let mut c = cfg(Provider::OpenAi);
|
||
c.openai_api = OpenAiApi::Responses;
|
||
c
|
||
}
|
||
|
||
fn tool_call_history() -> Vec<HistoryItem> {
|
||
vec![
|
||
HistoryItem::User("call the tool".into()),
|
||
HistoryItem::Assistant {
|
||
text: "ok, calling".into(),
|
||
tool_calls: vec![ToolCall {
|
||
provider_id: "call_abc".into(),
|
||
name: "dev__shell".into(),
|
||
arguments: serde_json::json!({"command": "ls"}),
|
||
}],
|
||
},
|
||
HistoryItem::ToolResult(ToolResult {
|
||
provider_id: "call_abc".into(),
|
||
content: vec![ToolResultContent::Text("file.txt".into())],
|
||
is_error: false,
|
||
}),
|
||
]
|
||
}
|
||
|
||
#[test]
|
||
fn responses_body_top_level_shape() {
|
||
let tools = vec![ToolDef {
|
||
name: "dev__shell".into(),
|
||
description: "run a shell command".into(),
|
||
input_schema: serde_json::json!({
|
||
"type": "object",
|
||
"properties": {"command": {"type": "string"}},
|
||
}),
|
||
}];
|
||
let body = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&tools,
|
||
"model",
|
||
None,
|
||
);
|
||
assert_eq!(body["model"], "model");
|
||
assert_eq!(body["instructions"], "system");
|
||
assert_eq!(body["max_output_tokens"], 1024);
|
||
assert!(
|
||
body.get("messages").is_none(),
|
||
"must use `input`, not `messages`"
|
||
);
|
||
assert!(body.get("max_tokens").is_none());
|
||
assert!(body.get("max_completion_tokens").is_none());
|
||
|
||
// Tools are flat — top-level type/name/description/parameters.
|
||
let tool = &body["tools"][0];
|
||
assert_eq!(tool["type"], "function");
|
||
assert_eq!(tool["name"], "dev__shell");
|
||
assert!(
|
||
tool.get("function").is_none(),
|
||
"Responses tool schema is flat"
|
||
);
|
||
assert_eq!(body["tool_choice"], "auto");
|
||
}
|
||
|
||
#[test]
|
||
fn responses_body_replay_emits_function_call_before_output() {
|
||
// Replay requirement from the live API: the assistant's prior
|
||
// function_call item *must* appear in `input[]` before its matching
|
||
// function_call_output, otherwise the API rejects with
|
||
// "No tool call found for call_id ...".
|
||
let body = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&tool_call_history(),
|
||
&[],
|
||
"model",
|
||
None,
|
||
);
|
||
let input = body["input"].as_array().unwrap();
|
||
|
||
// [0] user, [1] assistant text, [2] function_call, [3] function_call_output
|
||
assert_eq!(input[0]["role"], "user");
|
||
assert_eq!(input[0]["content"][0]["type"], "input_text");
|
||
assert_eq!(input[0]["content"][0]["text"], "call the tool");
|
||
|
||
assert_eq!(input[1]["role"], "assistant");
|
||
assert_eq!(input[1]["content"][0]["type"], "output_text");
|
||
assert_eq!(input[1]["content"][0]["text"], "ok, calling");
|
||
|
||
assert_eq!(input[2]["type"], "function_call");
|
||
assert_eq!(input[2]["call_id"], "call_abc");
|
||
assert_eq!(input[2]["name"], "dev__shell");
|
||
// Arguments are a JSON-encoded string per spec.
|
||
assert_eq!(input[2]["arguments"], "{\"command\":\"ls\"}");
|
||
|
||
assert_eq!(input[3]["type"], "function_call_output");
|
||
assert_eq!(input[3]["call_id"], "call_abc");
|
||
assert_eq!(input[3]["output"], "file.txt");
|
||
}
|
||
|
||
#[test]
|
||
fn responses_body_skips_empty_assistant_text() {
|
||
// Mirrors the Chat Completions behavior (#559/#560): empty assistant
|
||
// turns are skipped so we don't emit an empty `output_text` block,
|
||
// but the tool_call(s) on that assistant turn still go through.
|
||
let history = vec![
|
||
HistoryItem::User("u".into()),
|
||
HistoryItem::Assistant {
|
||
text: String::new(),
|
||
tool_calls: vec![ToolCall {
|
||
provider_id: "call_x".into(),
|
||
name: "t".into(),
|
||
arguments: serde_json::json!({}),
|
||
}],
|
||
},
|
||
];
|
||
let body = responses_body(&cfg_responses(), "system", &history, &[], "model", None);
|
||
let input = body["input"].as_array().unwrap();
|
||
assert_eq!(input.len(), 2);
|
||
assert_eq!(input[0]["role"], "user");
|
||
assert_eq!(input[1]["type"], "function_call");
|
||
}
|
||
|
||
#[test]
|
||
fn responses_body_image_tool_result_attaches_input_image() {
|
||
let body = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&image_history(),
|
||
&[],
|
||
"model",
|
||
None,
|
||
);
|
||
let input = body["input"].as_array().unwrap();
|
||
// function_call_output carries the text part; image rides on a
|
||
// trailing user message as `input_image`.
|
||
let fco = input
|
||
.iter()
|
||
.find(|i| i["type"] == "function_call_output")
|
||
.unwrap();
|
||
assert_eq!(fco["call_id"], "toolu_1");
|
||
let img_msg = input.iter().rev().find(|i| i["role"] == "user").unwrap();
|
||
assert_eq!(img_msg["content"][0]["type"], "input_image");
|
||
assert_eq!(
|
||
img_msg["content"][0]["image_url"],
|
||
"data:image/png;base64,aW1n"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn parse_responses_completed_with_text_is_end_turn() {
|
||
let v = serde_json::json!({
|
||
"status": "completed",
|
||
"output": [{
|
||
"type": "message",
|
||
"role": "assistant",
|
||
"content": [{"type": "output_text", "text": "hello"}],
|
||
}],
|
||
});
|
||
let r = parse_responses(v).unwrap();
|
||
assert_eq!(r.text, "hello");
|
||
assert!(r.tool_calls.is_empty());
|
||
assert_eq!(r.stop, ProviderStop::EndTurn);
|
||
}
|
||
|
||
#[test]
|
||
fn parse_responses_completed_with_function_call_is_tool_use() {
|
||
let v = serde_json::json!({
|
||
"status": "completed",
|
||
"output": [
|
||
{"type": "reasoning", "id": "rs_1", "summary": []},
|
||
{
|
||
"type": "function_call",
|
||
"call_id": "call_z",
|
||
"name": "dev__shell",
|
||
"arguments": "{\"command\":\"ls\"}",
|
||
},
|
||
],
|
||
});
|
||
let r = parse_responses(v).unwrap();
|
||
assert_eq!(r.text, "");
|
||
assert_eq!(r.tool_calls.len(), 1);
|
||
assert_eq!(r.tool_calls[0].provider_id, "call_z");
|
||
assert_eq!(r.tool_calls[0].name, "dev__shell");
|
||
assert_eq!(
|
||
r.tool_calls[0].arguments,
|
||
serde_json::json!({"command": "ls"})
|
||
);
|
||
assert_eq!(r.stop, ProviderStop::ToolUse);
|
||
}
|
||
|
||
#[test]
|
||
fn parse_responses_incomplete_max_output_tokens() {
|
||
let v = serde_json::json!({
|
||
"status": "incomplete",
|
||
"incomplete_details": {"reason": "max_output_tokens"},
|
||
"output": [],
|
||
});
|
||
let r = parse_responses(v).unwrap();
|
||
assert_eq!(r.stop, ProviderStop::MaxTokens);
|
||
}
|
||
|
||
#[test]
|
||
fn is_responses_required_error_matrix() {
|
||
for (body, want) in [
|
||
// Databricks GPT-5.5 (the actual case we observed).
|
||
("Function tools with reasoning_effort are not supported for gpt-5.5 in /v1/chat/completions. Please use /v1/responses instead.", true),
|
||
// Forward-compat: OpenAI saying the same thing in prose.
|
||
("This model requires the Responses API. Please use the Responses API instead.", true),
|
||
// Negatives — must NOT trigger on unrelated 4xx.
|
||
("{\"error\":\"invalid_api_key\"}", false),
|
||
("max_tokens is not supported with this model", false),
|
||
("", false),
|
||
] {
|
||
assert_eq!(is_responses_required_error(body), want, "body={body:?}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn databricks_v2_routes_by_model_family() {
|
||
for (model, route, path) in [
|
||
(
|
||
"databricks-gpt-5-5",
|
||
DatabricksV2Route::OpenAiResponses,
|
||
"/ai-gateway/openai/v1/responses",
|
||
),
|
||
(
|
||
"databricks-claude-opus-4-7",
|
||
DatabricksV2Route::AnthropicMessages,
|
||
"/ai-gateway/anthropic/v1/messages",
|
||
),
|
||
(
|
||
"custom-tool-model",
|
||
DatabricksV2Route::MlflowChatCompletions,
|
||
"/ai-gateway/mlflow/v1/chat/completions",
|
||
),
|
||
] {
|
||
let got = databricks_v2_route_for_model(model);
|
||
assert_eq!(got, route, "model={model}");
|
||
assert_eq!(databricks_v2_path(got), path, "model={model}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn parse_responses_rejects_malformed_function_arguments() {
|
||
let v = serde_json::json!({
|
||
"status": "completed",
|
||
"output": [{
|
||
"type": "function_call",
|
||
"call_id": "call_z",
|
||
"name": "t",
|
||
"arguments": "not json {",
|
||
}],
|
||
});
|
||
assert!(matches!(parse_responses(v), Err(AgentError::Llm(_))));
|
||
}
|
||
|
||
#[test]
|
||
fn openai_tool_result_adds_followup_image_user_message() {
|
||
let body = openai_body(
|
||
&cfg(Provider::OpenAi),
|
||
"system",
|
||
&image_history(),
|
||
&[],
|
||
"model",
|
||
None,
|
||
);
|
||
assert_eq!(body["messages"][3]["role"], "tool");
|
||
assert!(body["messages"][3]["content"]
|
||
.as_str()
|
||
.unwrap()
|
||
.contains("provided in the next user message"));
|
||
assert_eq!(body["messages"][4]["role"], "user");
|
||
assert_eq!(body["messages"][4]["content"][0]["type"], "image_url");
|
||
assert_eq!(
|
||
body["messages"][4]["content"][0]["image_url"]["url"],
|
||
"data:image/png;base64,aW1n"
|
||
);
|
||
}
|
||
|
||
/// Regression for Databricks model serving (and any OpenAI-Chat frontend
|
||
/// that translates to Anthropic on the way to the model). Parallel tool
|
||
/// calls where one or more return images previously produced an
|
||
/// interleaved sequence:
|
||
/// role:"tool" (A)
|
||
/// role:"user" (image A)
|
||
/// role:"tool" (B)
|
||
/// role:"user" (image B)
|
||
/// The intervening user message split the run of tool results, so the
|
||
/// translator could no longer fold them into a single Anthropic
|
||
/// `tool_result`-bearing user message — Anthropic then rejected the
|
||
/// request with "tool_use ids were found without tool_result blocks
|
||
/// immediately after". Fix: every `role:"tool"` for a run of adjacent
|
||
/// ToolResults emits contiguously, then a single trailing user message
|
||
/// carries all of the images from the batch.
|
||
#[test]
|
||
fn openai_parallel_image_tool_results_stay_contiguous() {
|
||
let history = vec![
|
||
HistoryItem::User("describe both images".into()),
|
||
HistoryItem::Assistant {
|
||
text: String::new(),
|
||
tool_calls: vec![
|
||
ToolCall {
|
||
provider_id: "toolu_a".into(),
|
||
name: "dev__view_image".into(),
|
||
arguments: serde_json::json!({"source": "a.png"}),
|
||
},
|
||
ToolCall {
|
||
provider_id: "toolu_b".into(),
|
||
name: "dev__view_image".into(),
|
||
arguments: serde_json::json!({"source": "b.png"}),
|
||
},
|
||
],
|
||
},
|
||
HistoryItem::ToolResult(ToolResult {
|
||
provider_id: "toolu_a".into(),
|
||
content: vec![
|
||
ToolResultContent::Text("10×10, 70 B (image/png from a.png)".into()),
|
||
ToolResultContent::Image {
|
||
data: "aaa".into(),
|
||
mime_type: "image/png".into(),
|
||
},
|
||
],
|
||
is_error: false,
|
||
}),
|
||
HistoryItem::ToolResult(ToolResult {
|
||
provider_id: "toolu_b".into(),
|
||
content: vec![
|
||
ToolResultContent::Text("10×10, 70 B (image/png from b.png)".into()),
|
||
ToolResultContent::Image {
|
||
data: "bbb".into(),
|
||
mime_type: "image/png".into(),
|
||
},
|
||
],
|
||
is_error: false,
|
||
}),
|
||
];
|
||
let body = openai_body(
|
||
&cfg(Provider::OpenAi),
|
||
"system",
|
||
&history,
|
||
&[],
|
||
"model",
|
||
None,
|
||
);
|
||
let messages = body["messages"].as_array().unwrap();
|
||
// [0] system, [1] user, [2] assistant(tool_calls), [3] tool A, [4] tool B, [5] user(images)
|
||
assert_eq!(messages.len(), 6, "messages: {messages:#?}");
|
||
assert_eq!(messages[3]["role"], "tool");
|
||
assert_eq!(messages[3]["tool_call_id"], "toolu_a");
|
||
assert_eq!(
|
||
messages[4]["role"], "tool",
|
||
"tool results must stay adjacent; intervening user message breaks Databricks/Anthropic pairing"
|
||
);
|
||
assert_eq!(messages[4]["tool_call_id"], "toolu_b");
|
||
assert_eq!(messages[5]["role"], "user");
|
||
let imgs = messages[5]["content"].as_array().unwrap();
|
||
assert_eq!(imgs.len(), 2);
|
||
assert_eq!(imgs[0]["image_url"]["url"], "data:image/png;base64,aaa");
|
||
assert_eq!(imgs[1]["image_url"]["url"], "data:image/png;base64,bbb");
|
||
}
|
||
|
||
// ---- ThinkingEffort body-shape tests ----
|
||
|
||
#[test]
|
||
fn anthropic_body_omits_thinking_when_effort_none() {
|
||
let body = anthropic_body(
|
||
&cfg(Provider::Anthropic),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"model",
|
||
None,
|
||
);
|
||
assert!(
|
||
body.get("thinking").is_none(),
|
||
"thinking must be absent when effort is None"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_body_emits_thinking_when_effort_high() {
|
||
// claude-3.x model → manual budget_tokens shape.
|
||
// Use max_output_tokens = 4096 so budget fits: headroom = 4096 - 1024 = 3072.
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 4096;
|
||
let body = anthropic_body(
|
||
&c,
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"claude-3-7-sonnet-20250219",
|
||
Some(ThinkingEffort::High),
|
||
);
|
||
assert_eq!(body["thinking"]["type"], "enabled");
|
||
// budget_tokens = min(32768, 4096-1024) = 3072
|
||
assert_eq!(body["thinking"]["budget_tokens"], 3072);
|
||
assert!(body.get("output_config").is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_body_omits_thinking_when_max_output_too_small() {
|
||
// max_output_tokens = 2047: headroom = 2047 - 1024 = 1023 < 1024 → omit thinking.
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 2047;
|
||
let body = anthropic_body(
|
||
&c,
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"claude-3-7-sonnet-20250219",
|
||
Some(ThinkingEffort::High),
|
||
);
|
||
assert!(
|
||
body.get("thinking").is_none(),
|
||
"thinking must be omitted when max_output_tokens leaves < 1024 for budget"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_body_emits_thinking_at_boundary_2048() {
|
||
// max_output_tokens = 2048: headroom = 2048 - 1024 = 1024 ≥ 1024 → emit.
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 2048;
|
||
let body = anthropic_body(
|
||
&c,
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"claude-3-7-sonnet-20250219",
|
||
Some(ThinkingEffort::High),
|
||
);
|
||
let t = body
|
||
.get("thinking")
|
||
.expect("thinking must be present at boundary 2048");
|
||
assert_eq!(t["budget_tokens"], 1024); // min(32768, 2048-1024)
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_body_emits_thinking_high_uncapped_when_budget_fits() {
|
||
// When max_output_tokens is large enough, budget_tokens is not capped.
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 65_536;
|
||
let body = anthropic_body(
|
||
&c,
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"claude-3-7-sonnet-20250219",
|
||
Some(ThinkingEffort::High),
|
||
);
|
||
assert_eq!(body["thinking"]["budget_tokens"], 32_768);
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_body_emits_thinking_low_budget() {
|
||
// Low budget (1024 tokens) exactly fits when max_output_tokens = 2048.
|
||
// headroom = 2048 - 1024 = 1024; min(1024, 1024) = 1024 ≥ 1024 → emit.
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 2048;
|
||
let body = anthropic_body(
|
||
&c,
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"claude-3-7-sonnet-20250219",
|
||
Some(ThinkingEffort::Low),
|
||
);
|
||
// Low budget (1024) fits exactly at the boundary — emitted without capping.
|
||
assert_eq!(body["thinking"]["budget_tokens"], 1024);
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_body_emits_adaptive_thinking_for_opus_4() {
|
||
// Adaptive Claude (claude-opus-4-6/4.7/4.8) → thinking:{type:"adaptive"} + output_config.effort.
|
||
// Note: Opus 4.5 is NOT adaptive — it uses manual budget.
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 32_768;
|
||
let body = anthropic_body(
|
||
&c,
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"claude-opus-4-7",
|
||
Some(ThinkingEffort::High),
|
||
);
|
||
assert_eq!(
|
||
body["thinking"]["type"], "adaptive",
|
||
"thinking must be {{type:adaptive}} for claude-opus-4-7"
|
||
);
|
||
assert_eq!(body["output_config"]["effort"], "high");
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_body_emits_manual_budget_for_opus_4_5() {
|
||
// Opus 4.5 uses manual budget (effort page: "uses manual thinking").
|
||
// max_output_tokens = 32768; headroom = 32768 - 1024 = 31744; min(32768, 31744) = 31744.
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 32_768;
|
||
let body = anthropic_body(
|
||
&c,
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"claude-opus-4-5",
|
||
Some(ThinkingEffort::High),
|
||
);
|
||
assert_eq!(body["thinking"]["type"], "enabled");
|
||
assert_eq!(body["thinking"]["budget_tokens"], 31_744); // min(32768, 32768-1024)
|
||
assert!(
|
||
body.get("output_config").is_none(),
|
||
"output_config must be absent for claude-opus-4-5 (manual budget)"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_body_omits_both_fields_for_unrecognized_model() {
|
||
// Non-Anthropic models (gpt-5, llama, etc.) → omit both fields rather than guess.
|
||
let body = anthropic_body(
|
||
&cfg(Provider::Anthropic),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"gpt-4o",
|
||
Some(ThinkingEffort::High),
|
||
);
|
||
assert!(body.get("thinking").is_none(), "thinking must be absent");
|
||
assert!(
|
||
body.get("output_config").is_none(),
|
||
"output_config must be absent"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn openai_body_omits_reasoning_effort_when_none() {
|
||
let body = openai_body(
|
||
&cfg(Provider::OpenAi),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"model",
|
||
None,
|
||
);
|
||
assert!(
|
||
body.get("reasoning_effort").is_none(),
|
||
"reasoning_effort must be absent when effort is None"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn openai_body_emits_reasoning_effort_medium() {
|
||
let body = openai_body(
|
||
&cfg(Provider::OpenAi),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"model",
|
||
Some(ThinkingEffort::Medium),
|
||
);
|
||
assert_eq!(body["reasoning_effort"], "medium");
|
||
}
|
||
|
||
#[test]
|
||
fn responses_body_omits_reasoning_when_effort_none() {
|
||
let body = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"model",
|
||
None,
|
||
);
|
||
assert!(
|
||
body.get("reasoning").is_none(),
|
||
"reasoning must be absent when effort is None"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn responses_body_emits_reasoning_effort_low() {
|
||
let body = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"model",
|
||
Some(ThinkingEffort::Low),
|
||
);
|
||
assert_eq!(body["reasoning"]["effort"], "low");
|
||
}
|
||
|
||
#[test]
|
||
fn effective_model_overrides_cfg_model_in_anthropic_body() {
|
||
let body = anthropic_body(
|
||
&cfg(Provider::Anthropic),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"override-model",
|
||
None,
|
||
);
|
||
assert_eq!(body["model"], "override-model");
|
||
}
|
||
|
||
#[test]
|
||
fn effective_model_overrides_cfg_model_in_openai_body() {
|
||
let body = openai_body(
|
||
&cfg(Provider::OpenAi),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"override-model",
|
||
None,
|
||
);
|
||
assert_eq!(body["model"], "override-model");
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_body_opus_4_8_xhigh_emits_xhigh_effort() {
|
||
// Body-shape regression: xhigh on Opus 4.8 must emit output_config.effort="xhigh".
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 32_768;
|
||
let body = anthropic_body(
|
||
&c,
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"claude-opus-4-8",
|
||
Some(ThinkingEffort::XHigh),
|
||
);
|
||
assert_eq!(body["thinking"]["type"], "adaptive");
|
||
assert_eq!(body["output_config"]["effort"], "xhigh");
|
||
}
|
||
|
||
#[test]
|
||
fn anthropic_body_opus_4_8_max_emits_max_effort() {
|
||
// Body-shape regression: max on Opus 4.8 must emit output_config.effort="max".
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 32_768;
|
||
let body = anthropic_body(
|
||
&c,
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"claude-opus-4-8",
|
||
Some(ThinkingEffort::Max),
|
||
);
|
||
assert_eq!(body["thinking"]["type"], "adaptive");
|
||
assert_eq!(body["output_config"]["effort"], "max");
|
||
}
|
||
|
||
#[test]
|
||
fn openai_body_emits_xhigh_effort() {
|
||
// xhigh is a valid OpenAI effort value — must pass through.
|
||
let body = openai_body(
|
||
&cfg(Provider::OpenAi),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"model",
|
||
Some(ThinkingEffort::XHigh),
|
||
);
|
||
assert_eq!(body["reasoning_effort"], "xhigh");
|
||
}
|
||
|
||
#[test]
|
||
fn openai_body_emits_none_effort() {
|
||
// none is a valid OpenAI effort value.
|
||
let body = openai_body(
|
||
&cfg(Provider::OpenAi),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"model",
|
||
Some(ThinkingEffort::None),
|
||
);
|
||
assert_eq!(body["reasoning_effort"], "none");
|
||
}
|
||
|
||
#[test]
|
||
fn responses_body_emits_xhigh_effort() {
|
||
// xhigh is a valid Responses API effort value.
|
||
let body = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"model",
|
||
Some(ThinkingEffort::XHigh),
|
||
);
|
||
assert_eq!(body["reasoning"]["effort"], "xhigh");
|
||
}
|
||
|
||
#[test]
|
||
fn responses_body_emits_minimal_effort() {
|
||
// minimal is a valid Responses API effort value.
|
||
let body = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"model",
|
||
Some(ThinkingEffort::Minimal),
|
||
);
|
||
assert_eq!(body["reasoning"]["effort"], "minimal");
|
||
}
|
||
|
||
// ---- DatabricksV2 route-aware effort normalization (body-level assertions) ----
|
||
//
|
||
// The DBv2 `complete()` dispatch applies `normalize_effort_for_openai_route` /
|
||
// `normalize_effort_for_anthropic_route` before calling body builders. These tests
|
||
// verify the body shape that results from the already-normalized effort values — i.e.,
|
||
// they confirm the body builders correctly serialize the values the dispatch passes them.
|
||
|
||
#[test]
|
||
fn dbv2_openai_route_max_effort_clamped_to_xhigh_in_responses_body() {
|
||
// DBv2 GPT-5.5 route: max → clamped to xhigh by normalize_effort_for_openai_route
|
||
// before reaching responses_body. gpt-5.5 supports xhigh so the final value is xhigh.
|
||
let clamped =
|
||
crate::config::normalize_effort_for_openai_route(ThinkingEffort::Max, "gpt-5.5");
|
||
let body = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"gpt-5.5",
|
||
Some(clamped),
|
||
);
|
||
assert_eq!(
|
||
body["reasoning"]["effort"], "xhigh",
|
||
"DBv2 GPT-5.5 route: max must be clamped to xhigh before responses_body"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn dbv2_openai_route_max_effort_passes_through_for_gpt5_6() {
|
||
let normalized =
|
||
crate::config::normalize_effort_for_openai_route(ThinkingEffort::Max, "gpt-5.6-sol");
|
||
let body = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"gpt-5.6-sol",
|
||
Some(normalized),
|
||
);
|
||
assert_eq!(
|
||
body["reasoning"]["effort"], "max",
|
||
"DBv2 GPT-5.6 route must serialize max to the Responses API"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn dbv2_mlflow_route_max_effort_clamped_to_xhigh_in_openai_body() {
|
||
// DBv2 MLflow route (unknown model): max → clamped to xhigh by normalize_effort_for_openai_route.
|
||
// Unknown models pass through after the max→xhigh clamp.
|
||
let clamped =
|
||
crate::config::normalize_effort_for_openai_route(ThinkingEffort::Max, "llama-4");
|
||
let body = openai_body(
|
||
&cfg(Provider::OpenAi),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"llama-4",
|
||
Some(clamped),
|
||
);
|
||
assert_eq!(
|
||
body["reasoning_effort"], "xhigh",
|
||
"DBv2 MLflow route: max must be clamped to xhigh before openai_body"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn dbv2_openai_route_none_minimal_pass_through_in_responses_body() {
|
||
// Verify that supported values pass through for the respective model families.
|
||
// gpt-5.5 supports none (but not minimal); gpt-5 base supports minimal (but not none).
|
||
let none_normalized =
|
||
crate::config::normalize_effort_for_openai_route(ThinkingEffort::None, "gpt-5.5");
|
||
assert_eq!(
|
||
none_normalized,
|
||
ThinkingEffort::None,
|
||
"OpenAI normalizer must not touch none for gpt-5.5"
|
||
);
|
||
let minimal_normalized =
|
||
crate::config::normalize_effort_for_openai_route(ThinkingEffort::Minimal, "gpt-5");
|
||
assert_eq!(
|
||
minimal_normalized,
|
||
ThinkingEffort::Minimal,
|
||
"OpenAI normalizer must not touch minimal for gpt-5 base"
|
||
);
|
||
let body = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"gpt-5.5",
|
||
Some(none_normalized),
|
||
);
|
||
assert_eq!(
|
||
body["reasoning"]["effort"], "none",
|
||
"DBv2 GPT-5.5 route: none must be emitted as-is"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn dbv2_claude_route_none_effort_omits_thinking_fields() {
|
||
// DBv2 Claude route: none → normalize_effort_for_anthropic_route returns None → omit.
|
||
let normalized = crate::config::normalize_effort_for_anthropic_route(ThinkingEffort::None);
|
||
assert_eq!(
|
||
normalized, None,
|
||
"Anthropic normalizer must return None for ThinkingEffort::None"
|
||
);
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 32_768;
|
||
let body = anthropic_body(
|
||
&c,
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"claude-opus-4-8",
|
||
normalized, // None → omit thinking fields
|
||
);
|
||
assert!(
|
||
body.get("thinking").is_none(),
|
||
"DBv2 Claude route: none effort must omit thinking fields"
|
||
);
|
||
assert!(
|
||
body.get("output_config").is_none(),
|
||
"DBv2 Claude route: none effort must omit output_config"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn dbv2_route_switch_max_body_level_simulation() {
|
||
// Body-level simulation of a session/set_model switch from a Claude model to a GPT-5
|
||
// model when thinking_effort=max. Calls body builders and normalizers directly (not
|
||
// through the ACP session/set_model path or DatabricksV2 dispatch) to verify the
|
||
// correct output shape for each side of the route switch.
|
||
// Before the switch: Claude route → max passes through as Anthropic "max".
|
||
// After the switch: GPT-5 route → max clamped to xhigh.
|
||
let mut c = cfg(Provider::Anthropic);
|
||
c.max_output_tokens = 32_768;
|
||
|
||
// Before switch: claude-opus-4-8 with effort=max → adaptive shape, effort="max"
|
||
let (thinking_before, oc_before) = crate::config::anthropic_thinking_config(
|
||
"claude-opus-4-8",
|
||
ThinkingEffort::Max,
|
||
32_768,
|
||
);
|
||
assert_eq!(thinking_before.unwrap()["type"], "adaptive");
|
||
assert_eq!(oc_before.unwrap()["effort"], "max");
|
||
|
||
// After switch to GPT-5.5 route: normalize max → xhigh for responses_body
|
||
// (gpt-5.5 supports xhigh, so the clamp result is xhigh, not further reduced)
|
||
let clamped =
|
||
crate::config::normalize_effort_for_openai_route(ThinkingEffort::Max, "gpt-5.5");
|
||
assert_eq!(clamped, ThinkingEffort::XHigh);
|
||
let body_after = responses_body(
|
||
&cfg_responses(),
|
||
"system",
|
||
&[HistoryItem::User("hi".into())],
|
||
&[],
|
||
"gpt-5.5",
|
||
Some(clamped),
|
||
);
|
||
assert_eq!(
|
||
body_after["reasoning"]["effort"], "xhigh",
|
||
"After set_model to GPT-5.5: max must be clamped to xhigh"
|
||
);
|
||
}
|
||
|
||
/// Regression: a connection that is accepted and then dropped before any
|
||
/// HTTP response bytes are written surfaces as a reqwest request-class
|
||
/// error (not `is_connect()`, not `is_timeout()`). The retry predicate
|
||
/// must recognize it; otherwise transient TLS/h2/proxy hiccups bubble
|
||
/// out of the agent as `transport: error sending request ...`.
|
||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||
async fn post_retries_on_dropped_connection_before_response() {
|
||
use std::sync::atomic::{AtomicU32, Ordering};
|
||
use std::sync::Arc;
|
||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||
use tokio::net::TcpListener;
|
||
|
||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||
let url = format!("http://{}/v1/x", listener.local_addr().unwrap());
|
||
let accepts = Arc::new(AtomicU32::new(0));
|
||
let accepts_srv = accepts.clone();
|
||
|
||
tokio::spawn(async move {
|
||
loop {
|
||
let (mut sock, _) = match listener.accept().await {
|
||
Ok(p) => p,
|
||
Err(_) => return,
|
||
};
|
||
let n = accepts_srv.fetch_add(1, Ordering::SeqCst);
|
||
if n == 0 {
|
||
// First attempt: read the request, then drop the socket
|
||
// without writing a response. reqwest surfaces this as
|
||
// a request-class error (is_request() == true).
|
||
let mut tmp = [0u8; 4096];
|
||
let _ = sock.read(&mut tmp).await;
|
||
drop(sock);
|
||
continue;
|
||
}
|
||
// Subsequent attempts: serve a tiny JSON body.
|
||
let mut buf = Vec::new();
|
||
let mut tmp = [0u8; 4096];
|
||
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
|
||
match sock.read(&mut tmp).await {
|
||
Ok(0) | Err(_) => return,
|
||
Ok(k) => buf.extend_from_slice(&tmp[..k]),
|
||
}
|
||
}
|
||
let body = "{\"ok\":true}";
|
||
let resp = format!(
|
||
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
|
||
Content-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||
body.len(),
|
||
body,
|
||
);
|
||
let _ = sock.write_all(resp.as_bytes()).await;
|
||
let _ = sock.shutdown().await;
|
||
}
|
||
});
|
||
|
||
let client = Client::builder()
|
||
.timeout(Duration::from_secs(5))
|
||
.build()
|
||
.unwrap();
|
||
let out = post(&client, &url, &serde_json::json!({}), false, |b| b)
|
||
.await
|
||
.expect("post should succeed after retry");
|
||
assert_eq!(out, serde_json::json!({ "ok": true }));
|
||
assert!(
|
||
accepts.load(Ordering::SeqCst) >= 2,
|
||
"server should have seen at least 2 connection attempts, saw {}",
|
||
accepts.load(Ordering::SeqCst)
|
||
);
|
||
}
|
||
|
||
/// A 499 response is retried and the call succeeds on the second attempt.
|
||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||
async fn post_retries_499_and_succeeds() {
|
||
use std::sync::atomic::{AtomicU32, Ordering};
|
||
use std::sync::Arc;
|
||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||
use tokio::net::TcpListener;
|
||
|
||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||
let url = format!("http://{}/v1/x", listener.local_addr().unwrap());
|
||
let accepts = Arc::new(AtomicU32::new(0));
|
||
let accepts_srv = accepts.clone();
|
||
|
||
tokio::spawn(async move {
|
||
loop {
|
||
let (mut sock, _) = match listener.accept().await {
|
||
Ok(p) => p,
|
||
Err(_) => return,
|
||
};
|
||
let n = accepts_srv.fetch_add(1, Ordering::SeqCst);
|
||
// Read the full request headers before responding.
|
||
let mut buf = Vec::new();
|
||
let mut tmp = [0u8; 4096];
|
||
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
|
||
match sock.read(&mut tmp).await {
|
||
Ok(0) | Err(_) => return,
|
||
Ok(k) => buf.extend_from_slice(&tmp[..k]),
|
||
}
|
||
}
|
||
if n == 0 {
|
||
// First attempt: respond with 499.
|
||
let resp = "HTTP/1.1 499 Client Closed Request\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
|
||
let _ = sock.write_all(resp.as_bytes()).await;
|
||
let _ = sock.shutdown().await;
|
||
continue;
|
||
}
|
||
// Subsequent attempts: 200 OK with a tiny JSON body.
|
||
let body = "{\"ok\":true}";
|
||
let resp = format!(
|
||
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
|
||
Content-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||
body.len(),
|
||
body,
|
||
);
|
||
let _ = sock.write_all(resp.as_bytes()).await;
|
||
let _ = sock.shutdown().await;
|
||
}
|
||
});
|
||
|
||
let client = Client::builder()
|
||
.timeout(Duration::from_secs(5))
|
||
.build()
|
||
.unwrap();
|
||
let out = post(&client, &url, &serde_json::json!({}), false, |b| b)
|
||
.await
|
||
.expect("post should succeed after 499 retry");
|
||
assert_eq!(out, serde_json::json!({ "ok": true }));
|
||
assert!(
|
||
accepts.load(Ordering::SeqCst) >= 2,
|
||
"server must see at least 2 attempts (got {})",
|
||
accepts.load(Ordering::SeqCst)
|
||
);
|
||
}
|
||
|
||
/// When all MAX_RETRIES attempts return 499 the error includes "exhausted retries".
|
||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||
async fn post_exhausts_retries_on_persistent_499() {
|
||
use std::sync::atomic::{AtomicU32, Ordering};
|
||
use std::sync::Arc;
|
||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||
use tokio::net::TcpListener;
|
||
|
||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||
let url = format!("http://{}/v1/x", listener.local_addr().unwrap());
|
||
let accepts = Arc::new(AtomicU32::new(0));
|
||
let accepts_srv = accepts.clone();
|
||
|
||
tokio::spawn(async move {
|
||
loop {
|
||
let (mut sock, _) = match listener.accept().await {
|
||
Ok(p) => p,
|
||
Err(_) => return,
|
||
};
|
||
accepts_srv.fetch_add(1, Ordering::SeqCst);
|
||
// Read the full request headers before responding.
|
||
let mut buf = Vec::new();
|
||
let mut tmp = [0u8; 4096];
|
||
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
|
||
match sock.read(&mut tmp).await {
|
||
Ok(0) | Err(_) => return,
|
||
Ok(k) => buf.extend_from_slice(&tmp[..k]),
|
||
}
|
||
}
|
||
// Always 499.
|
||
let resp = "HTTP/1.1 499 Client Closed Request\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
|
||
let _ = sock.write_all(resp.as_bytes()).await;
|
||
let _ = sock.shutdown().await;
|
||
}
|
||
});
|
||
|
||
let client = Client::builder()
|
||
.timeout(Duration::from_secs(5))
|
||
.build()
|
||
.unwrap();
|
||
let err = post(&client, &url, &serde_json::json!({}), false, |b| b)
|
||
.await
|
||
.unwrap_err();
|
||
match &err {
|
||
PostError::Agent(AgentError::Llm(msg)) => {
|
||
assert!(
|
||
msg.contains("exhausted retries") && msg.contains("499"),
|
||
"expected 'exhausted retries' + '499' in error, got: {msg}"
|
||
);
|
||
assert!(
|
||
msg.contains("cumulative") && msg.contains("3 attempts"),
|
||
"expected cumulative duration + exact attempt count, got: {msg}"
|
||
);
|
||
}
|
||
other => panic!("expected PostError::Agent(AgentError::Llm), got: {other:?}"),
|
||
}
|
||
assert_eq!(
|
||
accepts.load(Ordering::SeqCst),
|
||
MAX_RETRIES,
|
||
"server must see exactly MAX_RETRIES attempts — 499 must be retried"
|
||
);
|
||
}
|
||
|
||
/// `terminal_llm_error` below `STALL_NOTICE_THRESHOLD` carries the detail
|
||
/// and attempt count but no stall-specific text — this is the common case
|
||
/// (a handful of quick retries), not an outage.
|
||
#[test]
|
||
fn terminal_llm_error_below_threshold_carries_detail_no_stall_claim() {
|
||
let err = terminal_llm_error(Duration::from_secs(2), 3, "exhausted retries: 499: body");
|
||
match err {
|
||
AgentError::Llm(msg) => {
|
||
assert!(msg.contains("cumulative 2s"), "must carry duration: {msg}");
|
||
assert!(
|
||
msg.contains("3 attempts"),
|
||
"must carry attempt count: {msg}"
|
||
);
|
||
assert!(
|
||
msg.contains("exhausted retries") && msg.contains("499"),
|
||
"must preserve detail: {msg}"
|
||
);
|
||
}
|
||
other => panic!("expected Llm, got: {other:?}"),
|
||
}
|
||
}
|
||
|
||
/// Above `STALL_NOTICE_THRESHOLD`, the error still carries the same
|
||
/// detail/duration/count — the threshold only gates the extra
|
||
/// `tracing::warn!` (not separately assertable here), not the error text.
|
||
/// Singular "1 attempt" (not "1 attempts") confirms the pluralization.
|
||
#[test]
|
||
fn terminal_llm_error_above_threshold_uses_singular_attempt() {
|
||
let err = terminal_llm_error(Duration::from_secs(301), 1, "transport: connection reset");
|
||
match err {
|
||
AgentError::Llm(msg) => {
|
||
assert!(
|
||
msg.contains("cumulative 301s"),
|
||
"must carry duration: {msg}"
|
||
);
|
||
assert!(msg.contains("1 attempt"), "must carry count: {msg}");
|
||
assert!(
|
||
!msg.contains("1 attempts"),
|
||
"singular for count of 1: {msg}"
|
||
);
|
||
}
|
||
other => panic!("expected Llm, got: {other:?}"),
|
||
}
|
||
}
|
||
|
||
/// Captures `tracing::warn!` events emitted during a scoped subscriber
|
||
/// so a test can assert on the stall-notice threshold without a real
|
||
/// sleep or a global logger.
|
||
struct StallWarnCapture {
|
||
count: Arc<std::sync::atomic::AtomicUsize>,
|
||
}
|
||
|
||
struct StallWarnVisitor {
|
||
saw_cumulative_stall: bool,
|
||
saw_attempts: bool,
|
||
}
|
||
|
||
impl tracing::field::Visit for StallWarnVisitor {
|
||
fn record_debug(&mut self, field: &tracing::field::Field, _value: &dyn std::fmt::Debug) {
|
||
match field.name() {
|
||
"cumulative_stall" => self.saw_cumulative_stall = true,
|
||
"attempts" => self.saw_attempts = true,
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
fn record_u64(&mut self, field: &tracing::field::Field, _value: u64) {
|
||
if field.name() == "attempts" {
|
||
self.saw_attempts = true;
|
||
}
|
||
}
|
||
}
|
||
|
||
impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for StallWarnCapture {
|
||
fn on_event(
|
||
&self,
|
||
event: &tracing::Event<'_>,
|
||
_ctx: tracing_subscriber::layer::Context<'_, S>,
|
||
) {
|
||
if *event.metadata().level() != tracing::Level::WARN {
|
||
return;
|
||
}
|
||
let mut visitor = StallWarnVisitor {
|
||
saw_cumulative_stall: false,
|
||
saw_attempts: false,
|
||
};
|
||
event.record(&mut visitor);
|
||
if visitor.saw_cumulative_stall && visitor.saw_attempts {
|
||
self.count.fetch_add(1, Ordering::SeqCst);
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Runs `f` under a scoped subscriber that only counts `WARN` events
|
||
/// carrying both `cumulative_stall` and `attempts` fields — the shape
|
||
/// `terminal_llm_error`'s stall notice emits. Returns that count.
|
||
fn count_stall_warnings(f: impl FnOnce()) -> usize {
|
||
let count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||
let subscriber = tracing_subscriber::registry().with(StallWarnCapture {
|
||
count: count.clone(),
|
||
});
|
||
tracing::subscriber::with_default(subscriber, f);
|
||
count.load(Ordering::SeqCst)
|
||
}
|
||
|
||
/// Below `STALL_NOTICE_THRESHOLD`, `terminal_llm_error` must not emit the
|
||
/// stall warning at all — otherwise every ordinary retry-exhaustion would
|
||
/// falsely read as an outage. This is mutation-sensitive: deleting the
|
||
/// threshold branch, weakening `>=` to `>` at the boundary, or an
|
||
/// always-warn mutation are all caught by pairing this with the 300s case
|
||
/// below.
|
||
#[test]
|
||
fn terminal_llm_error_below_threshold_emits_no_stall_warning() {
|
||
let warnings = count_stall_warnings(|| {
|
||
let _ = terminal_llm_error(Duration::from_secs(299), 3, "exhausted retries: 499");
|
||
});
|
||
assert_eq!(warnings, 0, "no stall warning below STALL_NOTICE_THRESHOLD");
|
||
}
|
||
|
||
/// At `STALL_NOTICE_THRESHOLD`, `terminal_llm_error` must emit exactly
|
||
/// one stall warning carrying the duration and attempt count — the
|
||
/// never-warn mutation is caught here; combined with the 299s case above,
|
||
/// a deleted branch or an always-warn mutation are caught by whichever
|
||
/// assertion it violates.
|
||
#[test]
|
||
fn terminal_llm_error_at_threshold_emits_one_stall_warning() {
|
||
let warnings = count_stall_warnings(|| {
|
||
let _ = terminal_llm_error(Duration::from_secs(300), 5, "transport: connection reset");
|
||
});
|
||
assert_eq!(
|
||
warnings, 1,
|
||
"exactly one stall warning with duration+attempts fields at threshold"
|
||
);
|
||
}
|
||
|
||
// ---- usage / input-token extraction -------------------------------------
|
||
|
||
#[test]
|
||
fn parse_anthropic_sums_input_and_cache_tokens() {
|
||
// input_tokens alone excludes cached tokens; the inclusive total must
|
||
// sum all three so a cache-heavy turn can't undercount the budget.
|
||
let v = serde_json::json!({
|
||
"stop_reason": "end_turn",
|
||
"content": [{"type": "text", "text": "hi"}],
|
||
"usage": {
|
||
"input_tokens": 100,
|
||
"cache_read_input_tokens": 900,
|
||
"cache_creation_input_tokens": 50,
|
||
"output_tokens": 7
|
||
}
|
||
});
|
||
let r = parse_anthropic(v).unwrap();
|
||
assert_eq!(r.input_tokens, Some(1050));
|
||
}
|
||
|
||
#[test]
|
||
fn parse_anthropic_input_tokens_only() {
|
||
let v = serde_json::json!({
|
||
"stop_reason": "end_turn",
|
||
"content": [{"type": "text", "text": "hi"}],
|
||
"usage": {"input_tokens": 42, "output_tokens": 3}
|
||
});
|
||
assert_eq!(parse_anthropic(v).unwrap().input_tokens, Some(42));
|
||
}
|
||
|
||
#[test]
|
||
fn parse_anthropic_missing_usage_is_none() {
|
||
let v = serde_json::json!({
|
||
"stop_reason": "end_turn",
|
||
"content": [{"type": "text", "text": "hi"}]
|
||
});
|
||
assert_eq!(parse_anthropic(v).unwrap().input_tokens, None);
|
||
}
|
||
|
||
#[test]
|
||
fn parse_openai_uses_prompt_tokens() {
|
||
let v = serde_json::json!({
|
||
"choices": [{"finish_reason": "stop", "message": {"content": "hi"}}],
|
||
"usage": {"prompt_tokens": 123, "completion_tokens": 4, "total_tokens": 127}
|
||
});
|
||
assert_eq!(parse_openai(v).unwrap().input_tokens, Some(123));
|
||
}
|
||
|
||
#[test]
|
||
fn parse_openai_databricks_sums_cache_fields() {
|
||
// Databricks uses the OpenAI chat wire format (prompt_tokens) but also
|
||
// reports Anthropic-style cache fields; the inclusive total sums them.
|
||
let v = serde_json::json!({
|
||
"choices": [{"finish_reason": "stop", "message": {"content": "hi"}}],
|
||
"usage": {
|
||
"prompt_tokens": 200,
|
||
"completion_tokens": 4,
|
||
"total_tokens": 204,
|
||
"cache_read_input_tokens": 800,
|
||
"cache_creation_input_tokens": 0
|
||
}
|
||
});
|
||
assert_eq!(parse_openai(v).unwrap().input_tokens, Some(1000));
|
||
}
|
||
|
||
#[test]
|
||
fn parse_openai_missing_usage_is_none() {
|
||
let v = serde_json::json!({
|
||
"choices": [{"finish_reason": "stop", "message": {"content": "hi"}}]
|
||
});
|
||
assert_eq!(parse_openai(v).unwrap().input_tokens, None);
|
||
}
|
||
|
||
#[test]
|
||
fn parse_responses_uses_input_tokens() {
|
||
let v = serde_json::json!({
|
||
"status": "completed",
|
||
"output": [{
|
||
"type": "message",
|
||
"role": "assistant",
|
||
"content": [{"type": "output_text", "text": "hi"}]
|
||
}],
|
||
"usage": {"input_tokens": 321, "output_tokens": 9, "total_tokens": 330}
|
||
});
|
||
assert_eq!(parse_responses(v).unwrap().input_tokens, Some(321));
|
||
}
|
||
|
||
#[test]
|
||
fn parse_responses_missing_usage_is_none() {
|
||
let v = serde_json::json!({
|
||
"status": "completed",
|
||
"output": [{
|
||
"type": "message",
|
||
"role": "assistant",
|
||
"content": [{"type": "output_text", "text": "hi"}]
|
||
}]
|
||
});
|
||
assert_eq!(parse_responses(v).unwrap().input_tokens, None);
|
||
}
|
||
|
||
#[test]
|
||
fn sum_usage_empty_object_is_none() {
|
||
// A `usage` object present but carrying none of the requested fields
|
||
// is "no usable reading" -> None, not Some(0).
|
||
let v = serde_json::json!({"usage": {"output_tokens": 5}});
|
||
assert_eq!(sum_usage(&v, &["input_tokens", "prompt_tokens"]), None);
|
||
}
|
||
|
||
/// A token source whose `bearer()` always hands back the same stale
|
||
/// token and whose `refresh_now()` mints a distinct fresh one, counting
|
||
/// each refresh. Lets a test assert exactly how many forced refreshes a
|
||
/// `post_openai` call provoked.
|
||
struct CountingAuth {
|
||
refreshes: std::sync::atomic::AtomicU32,
|
||
}
|
||
|
||
#[async_trait::async_trait]
|
||
impl TokenSource for CountingAuth {
|
||
async fn bearer(&self) -> Result<String, AgentError> {
|
||
Ok("stale".into())
|
||
}
|
||
async fn refresh_now(&self, _rejected: &str) -> Result<String, AgentError> {
|
||
self.refreshes
|
||
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||
Ok("fresh".into())
|
||
}
|
||
}
|
||
|
||
/// Stub that answers `reject_status` to any request carrying `Bearer
|
||
/// stale` and 200 to `Bearer fresh`. When `always_reject` is set it rejects
|
||
/// unconditionally, simulating a token the refresh can never satisfy.
|
||
/// Counts requests so a test can assert "one retry, not a loop".
|
||
async fn spawn_auth_stub(
|
||
always_reject: std::sync::Arc<std::sync::atomic::AtomicBool>,
|
||
reject_status: u16,
|
||
) -> String {
|
||
use std::sync::atomic::Ordering;
|
||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||
use tokio::net::TcpListener;
|
||
|
||
let reject_line = match reject_status {
|
||
401 => "401 Unauthorized",
|
||
403 => "403 Forbidden",
|
||
other => panic!("unsupported reject_status {other}"),
|
||
};
|
||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||
tokio::spawn(async move {
|
||
loop {
|
||
let (mut sock, _) = match listener.accept().await {
|
||
Ok(p) => p,
|
||
Err(_) => return,
|
||
};
|
||
let always_reject = always_reject.clone();
|
||
tokio::spawn(async move {
|
||
let mut buf = Vec::new();
|
||
let mut tmp = [0u8; 4096];
|
||
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
|
||
match sock.read(&mut tmp).await {
|
||
Ok(0) | Err(_) => return,
|
||
Ok(k) => buf.extend_from_slice(&tmp[..k]),
|
||
}
|
||
}
|
||
let head = String::from_utf8_lossy(&buf).to_ascii_lowercase();
|
||
let stale = head.contains("authorization: bearer stale");
|
||
let resp = if always_reject.load(Ordering::SeqCst) || stale {
|
||
format!(
|
||
"HTTP/1.1 {reject_line}\r\nContent-Length: 11\r\n\
|
||
Connection: close\r\n\r\ntoken stale"
|
||
)
|
||
} else {
|
||
let body = "{\"ok\":true}";
|
||
format!(
|
||
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\
|
||
Content-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||
body.len(),
|
||
body,
|
||
)
|
||
};
|
||
let _ = sock.write_all(resp.as_bytes()).await;
|
||
let _ = sock.shutdown().await;
|
||
});
|
||
}
|
||
});
|
||
url
|
||
}
|
||
|
||
fn llm_with(auth: Arc<dyn TokenSource>) -> Llm {
|
||
Llm {
|
||
http: Client::builder()
|
||
.timeout(Duration::from_secs(5))
|
||
.build()
|
||
.unwrap(),
|
||
auto_upgraded: std::sync::atomic::AtomicBool::new(false),
|
||
mesh_auto_state: Mutex::new(MeshAutoState::default()),
|
||
auth,
|
||
}
|
||
}
|
||
|
||
/// A single 401 forces exactly one refresh, the retry with the fresh
|
||
/// token succeeds, and a *later* call gets its own refresh — proving the
|
||
/// one-shot guard is per-call, not stored on the source.
|
||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||
async fn post_openai_refreshes_once_per_call_on_401() {
|
||
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
|
||
|
||
let always_401 = Arc::new(AtomicBool::new(false));
|
||
let base = spawn_auth_stub(always_401, 401).await;
|
||
let auth = Arc::new(CountingAuth {
|
||
refreshes: AtomicU32::new(0),
|
||
});
|
||
let llm = llm_with(auth.clone());
|
||
let mut c = cfg(Provider::OpenAi);
|
||
c.base_url = base;
|
||
|
||
let out = llm
|
||
.post_openai(&c, "/v1/x", &json!({}), "model")
|
||
.await
|
||
.expect("retry with fresh token should succeed");
|
||
assert_eq!(out, json!({ "ok": true }));
|
||
assert_eq!(auth.refreshes.load(Ordering::SeqCst), 1, "one refresh");
|
||
|
||
// Second call's 401 must trigger its own refresh — the guard cannot
|
||
// be a stored flag that an earlier turn already tripped.
|
||
let out2 = llm
|
||
.post_openai(&c, "/v1/x", &json!({}), "model")
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(out2, json!({ "ok": true }));
|
||
assert_eq!(
|
||
auth.refreshes.load(Ordering::SeqCst),
|
||
2,
|
||
"later call gets its own retry"
|
||
);
|
||
}
|
||
|
||
/// A persistent 401 (even the refreshed token is rejected) propagates as
|
||
/// `LlmAuth` after exactly one refresh — no infinite loop.
|
||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||
async fn post_openai_persistent_401_propagates_after_one_retry() {
|
||
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
|
||
|
||
let always_401 = Arc::new(AtomicBool::new(true));
|
||
let base = spawn_auth_stub(always_401, 401).await;
|
||
let auth = Arc::new(CountingAuth {
|
||
refreshes: AtomicU32::new(0),
|
||
});
|
||
let llm = llm_with(auth.clone());
|
||
let mut c = cfg(Provider::OpenAi);
|
||
c.base_url = base;
|
||
|
||
let err = llm
|
||
.post_openai(&c, "/v1/x", &json!({}), "model")
|
||
.await
|
||
.unwrap_err();
|
||
assert!(
|
||
matches!(err, PostError::Agent(AgentError::LlmAuth(_))),
|
||
"got {err:?}"
|
||
);
|
||
assert_eq!(
|
||
auth.refreshes.load(Ordering::SeqCst),
|
||
1,
|
||
"exactly one refresh, then propagate"
|
||
);
|
||
}
|
||
|
||
/// A 403 is treated as refreshable: a persistent 403 forces exactly one
|
||
/// refresh-and-retry, then propagates as `LlmAuth`. Proves a revoked-token
|
||
/// 403 takes the same recovery path as a 401, bounded by the per-call guard.
|
||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||
async fn post_openai_persistent_403_propagates_after_one_retry() {
|
||
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
|
||
|
||
let always_403 = Arc::new(AtomicBool::new(true));
|
||
let base = spawn_auth_stub(always_403, 403).await;
|
||
let auth = Arc::new(CountingAuth {
|
||
refreshes: AtomicU32::new(0),
|
||
});
|
||
let llm = llm_with(auth.clone());
|
||
let mut c = cfg(Provider::OpenAi);
|
||
c.base_url = base;
|
||
|
||
let err = llm
|
||
.post_openai(&c, "/v1/x", &json!({}), "model")
|
||
.await
|
||
.unwrap_err();
|
||
assert!(
|
||
matches!(err, PostError::Agent(AgentError::LlmAuth(_))),
|
||
"got {err:?}"
|
||
);
|
||
assert_eq!(
|
||
auth.refreshes.load(Ordering::SeqCst),
|
||
1,
|
||
"403 refreshes exactly once, then propagates"
|
||
);
|
||
}
|
||
|
||
/// A recoverable 403 (stale token 403s, fresh token 200s) forces exactly
|
||
/// one refresh and the retry succeeds — proving a 403 enters the refresh
|
||
/// path and a refreshed token clears it, the stale-token-403 recovery case.
|
||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||
async fn post_openai_refreshes_once_on_403() {
|
||
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
|
||
|
||
let always_403 = Arc::new(AtomicBool::new(false));
|
||
let base = spawn_auth_stub(always_403, 403).await;
|
||
let auth = Arc::new(CountingAuth {
|
||
refreshes: AtomicU32::new(0),
|
||
});
|
||
let llm = llm_with(auth.clone());
|
||
let mut c = cfg(Provider::OpenAi);
|
||
c.base_url = base;
|
||
|
||
let out = llm
|
||
.post_openai(&c, "/v1/x", &json!({}), "model")
|
||
.await
|
||
.expect("retry with fresh token should clear the 403");
|
||
assert_eq!(out, json!({ "ok": true }));
|
||
assert_eq!(auth.refreshes.load(Ordering::SeqCst), 1, "one refresh");
|
||
}
|
||
|
||
/// The default `refresh_now()` on a static source returns the static
|
||
/// token unchanged — a key that can't refresh still answers harmlessly.
|
||
#[tokio::test]
|
||
async fn static_token_source_refresh_now_returns_static_token() {
|
||
let src = StaticTokenSource::new("static-key");
|
||
assert_eq!(src.refresh_now("rejected").await.unwrap(), "static-key");
|
||
}
|
||
|
||
// ── Output-token parsing tests ──────────────────────────────────────────
|
||
|
||
/// `parse_anthropic` extracts `output_tokens` from the usage object.
|
||
#[test]
|
||
fn parse_anthropic_output_tokens() {
|
||
let v = serde_json::json!({
|
||
"stop_reason": "end_turn",
|
||
"content": [{"type": "text", "text": "hi"}],
|
||
"usage": {"input_tokens": 42, "output_tokens": 7}
|
||
});
|
||
assert_eq!(parse_anthropic(v).unwrap().output_tokens, Some(7));
|
||
}
|
||
|
||
/// `parse_anthropic` returns `None` for `output_tokens` when usage is absent.
|
||
#[test]
|
||
fn parse_anthropic_output_tokens_missing_usage_is_none() {
|
||
let v = serde_json::json!({
|
||
"stop_reason": "end_turn",
|
||
"content": [{"type": "text", "text": "hi"}]
|
||
});
|
||
assert_eq!(parse_anthropic(v).unwrap().output_tokens, None);
|
||
}
|
||
|
||
/// `parse_openai` maps `completion_tokens` to `output_tokens`.
|
||
#[test]
|
||
fn parse_openai_output_tokens_from_completion_tokens() {
|
||
let v = serde_json::json!({
|
||
"choices": [{"finish_reason": "stop", "message": {"content": "hi"}}],
|
||
"usage": {"prompt_tokens": 123, "completion_tokens": 4, "total_tokens": 127}
|
||
});
|
||
assert_eq!(parse_openai(v).unwrap().output_tokens, Some(4));
|
||
}
|
||
|
||
/// `parse_openai` returns `None` for `output_tokens` when usage is absent.
|
||
#[test]
|
||
fn parse_openai_output_tokens_missing_usage_is_none() {
|
||
let v = serde_json::json!({
|
||
"choices": [{"finish_reason": "stop", "message": {"content": "hi"}}]
|
||
});
|
||
assert_eq!(parse_openai(v).unwrap().output_tokens, None);
|
||
}
|
||
|
||
/// `parse_responses` extracts `output_tokens` from the usage object.
|
||
#[test]
|
||
fn parse_responses_output_tokens() {
|
||
let v = serde_json::json!({
|
||
"status": "completed",
|
||
"output": [{
|
||
"type": "message",
|
||
"role": "assistant",
|
||
"content": [{"type": "output_text", "text": "hi"}]
|
||
}],
|
||
"usage": {"input_tokens": 321, "output_tokens": 9, "total_tokens": 330}
|
||
});
|
||
assert_eq!(parse_responses(v).unwrap().output_tokens, Some(9));
|
||
}
|
||
|
||
/// `parse_responses` returns `None` for `output_tokens` when usage is absent.
|
||
#[test]
|
||
fn parse_responses_output_tokens_missing_usage_is_none() {
|
||
let v = serde_json::json!({
|
||
"status": "completed",
|
||
"output": [{
|
||
"type": "message",
|
||
"role": "assistant",
|
||
"content": [{"type": "output_text", "text": "hi"}]
|
||
}]
|
||
});
|
||
assert_eq!(parse_responses(v).unwrap().output_tokens, None);
|
||
}
|
||
}
|