use anyhow::Result;
pub use crate::cache_profile::CacheProfile;
use futures::StreamExt;
use reqwest::{Client, RequestBuilder};
use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
use serde_json::{json, Map, Value};
use std::collections::{BTreeMap, HashMap};
use std::fmt;
use tokio::sync::mpsc;
// ── message types ─────────────────────────────────────────────────────────────
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct ToolCallFunction {
pub name: String,
pub arguments: String,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct ToolCall {
pub id: String,
#[serde(rename = "type")]
pub kind: String,
pub function: ToolCallFunction,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Role {
System,
User,
Assistant,
Tool,
Reflection,
Compaction,
Other(String),
}
impl Role {
pub fn as_str(&self) -> &str {
match self {
Self::System => "system",
Self::User => "user",
Self::Assistant => "assistant",
Self::Tool => "tool",
Self::Reflection => "reflection",
Self::Compaction => "compaction",
Self::Other(role) => role,
}
}
}
impl From<&str> for Role {
fn from(role: &str) -> Self {
match role {
"system" => Self::System,
"user" => Self::User,
"assistant" => Self::Assistant,
"tool" => Self::Tool,
"reflection" => Self::Reflection,
"compaction" => Self::Compaction,
other => Self::Other(other.to_string()),
}
}
}
impl From<String> for Role {
fn from(role: String) -> Self {
Self::from(role.as_str())
}
}
impl PartialEq<&str> for Role {
fn eq(&self, other: &&str) -> bool {
self.as_str() == *other
}
}
impl PartialEq<Role> for &str {
fn eq(&self, other: &Role) -> bool {
*self == other.as_str()
}
}
impl fmt::Display for Role {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
impl Serialize for Role {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(self.as_str())
}
}
impl<'de> Deserialize<'de> for Role {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let role = String::deserialize(deserializer)?;
if role.is_empty() {
return Err(de::Error::custom("message role cannot be empty"));
}
Ok(Self::from(role))
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Message {
pub role: Role,
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning_content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
}
impl Message {
pub fn system(s: impl Into<String>) -> Self {
Self {
role: Role::System,
content: Some(s.into()),
reasoning_content: None,
tool_calls: None,
tool_call_id: None,
}
}
pub fn user(s: impl Into<String>) -> Self {
Self {
role: Role::User,
content: Some(s.into()),
reasoning_content: None,
tool_calls: None,
tool_call_id: None,
}
}
pub fn assistant(s: impl Into<String>) -> Self {
Self {
role: Role::Assistant,
content: Some(s.into()),
reasoning_content: None,
tool_calls: None,
tool_call_id: None,
}
}
pub fn with_tool_calls(
calls: Vec<ToolCall>,
content: Option<String>,
reasoning_content: Option<String>,
) -> Self {
Self {
role: Role::Assistant,
content,
reasoning_content,
tool_calls: Some(calls),
tool_call_id: None,
}
}
pub fn tool_result(id: impl Into<String>, content: impl Into<String>) -> Self {
Self {
role: Role::Tool,
content: Some(content.into()),
reasoning_content: None,
tool_calls: None,
tool_call_id: Some(id.into()),
}
}
pub fn content_str(&self) -> &str {
self.content.as_deref().unwrap_or("")
}
fn normalize_for_api(&mut self) -> bool {
if self.role != Role::Assistant {
return true;
}
if matches!(self.tool_calls.as_ref(), Some(calls) if calls.is_empty()) {
self.tool_calls = None;
}
let has_tool_calls = self.tool_calls.is_some();
let has_content = self
.content
.as_deref()
.map(|content| !content.trim().is_empty())
.unwrap_or(false);
if !has_content && !has_tool_calls {
return false;
}
if has_tool_calls && self.content.is_none() {
self.content = Some(String::new());
}
true
}
}
// ── tool definitions ──────────────────────────────────────────────────────────
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ToolFunctionDef {
pub name: String,
pub description: String,
pub parameters: Value,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ToolDef {
#[serde(rename = "type")]
pub kind: String,
pub function: ToolFunctionDef,
}
impl ToolDef {
pub fn function(name: &str, description: &str, parameters: Value) -> Self {
Self {
kind: "function".into(),
function: ToolFunctionDef {
name: name.into(),
description: description.into(),
parameters,
},
}
}
}
// ── events ────────────────────────────────────────────────────────────────────
#[derive(Debug, Serialize, Clone, Default)]
pub struct Usage {
pub prompt_tokens: usize,
pub completion_tokens: usize,
pub total_tokens: usize,
/// `None` means the provider did not report cache usage; `Some(0)` is a
/// provider-reported zero, not an absent measurement.
#[serde(skip_serializing_if = "Option::is_none")]
pub cached_prompt_tokens: Option<usize>,
/// Tokens the route reported writing into a prompt cache.
///
/// Writes and reads are priced differently on routes that charge for them (a cached prefix
/// can cost more to write than to process uncached, and much less to read), so a reader that
/// only sees reads cannot price a prompt. `None` here means the route never said, which is a
/// different statement from a route that said zero.
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_write_tokens: Option<usize>,
/// Tokens the route reported as *not* reused, when it reports the miss side itself.
///
/// Reported, not derived: `prompt_tokens - cached_prompt_tokens` is our arithmetic, and a
/// route that counts its own cache in blocks can disagree with it.
#[serde(skip_serializing_if = "Option::is_none")]
pub uncached_prompt_tokens: Option<usize>,
/// Whether the provider actually reported token counts for this request.
///
/// `false` means the counts above are not measurements: the provider sent no
/// usage block and the stream fell back to defaults so the klbr-measured
/// timings could still be carried. Recording that as zero would invent a
/// measurement, and a cache miss would become indistinguishable from a
/// provider that never said anything.
pub tokens_reported: bool,
/// Time to first response byte measured by klbr, not provider cache usage.
pub prompt_processing_ms: u64,
pub generation_ms: u64,
pub bridge_processing_ms: u64,
}
/// One usage block, in every shape the routes we talk to actually send.
///
/// The shapes differ per route family and not per our choice: OpenAI reports
/// `prompt_tokens_details.cached_tokens` (and, on newer models, a write count under the same
/// details or `input_tokens_details`), DeepSeek reports `prompt_cache_hit_tokens` and
/// `prompt_cache_miss_tokens` next to `prompt_tokens`. Reading only one family's field names
/// makes the other family's cache usage invisible, and an invisible measurement gets read as a
/// zero by the next person who looks.
#[derive(Debug, Deserialize)]
struct UsageWire {
#[serde(default)]
prompt_tokens: Option<usize>,
#[serde(default)]
completion_tokens: Option<usize>,
#[serde(default)]
total_tokens: Option<usize>,
#[serde(default)]
cached_prompt_tokens: Option<usize>,
#[serde(default)]
prompt_cache_hit_tokens: Option<usize>,
#[serde(default)]
prompt_cache_miss_tokens: Option<usize>,
#[serde(default)]
cache_write_tokens: Option<usize>,
#[serde(default)]
prompt_tokens_details: Option<PromptTokensDetails>,
#[serde(default)]
input_tokens_details: Option<PromptTokensDetails>,
#[serde(default)]
prompt_processing_ms: Option<u64>,
#[serde(default)]
generation_ms: Option<u64>,
#[serde(default)]
bridge_processing_ms: Option<u64>,
}
#[derive(Debug, Deserialize)]
struct PromptTokensDetails {
#[serde(default)]
cached_tokens: Option<usize>,
#[serde(default)]
cache_write_tokens: Option<usize>,
}
impl<'de> Deserialize<'de> for Usage {
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let wire = UsageWire::deserialize(deserializer)?;
Ok(Self {
prompt_tokens: wire.prompt_tokens.unwrap_or_default(),
completion_tokens: wire.completion_tokens.unwrap_or_default(),
total_tokens: wire.total_tokens.unwrap_or_default(),
cached_prompt_tokens: wire
.cached_prompt_tokens
.or_else(|| {
wire.prompt_tokens_details
.as_ref()
.and_then(|details| details.cached_tokens)
})
.or(wire.prompt_cache_hit_tokens),
cache_write_tokens: wire.cache_write_tokens.or_else(|| {
wire.prompt_tokens_details
.as_ref()
.and_then(|details| details.cache_write_tokens)
.or_else(|| {
wire.input_tokens_details
.as_ref()
.and_then(|details| details.cache_write_tokens)
})
}),
uncached_prompt_tokens: wire.prompt_cache_miss_tokens,
// Reaching this deserializer means a provider usage object existed:
// this is the one place that can say "reported".
tokens_reported: true,
prompt_processing_ms: wire.prompt_processing_ms.unwrap_or_default(),
generation_ms: wire.generation_ms.unwrap_or_default(),
bridge_processing_ms: wire.bridge_processing_ms.unwrap_or_default(),
})
}
}
#[derive(Debug, Clone)]
pub enum LlmEvent {
Token(String),
ThinkToken(String),
Usage(Usage),
/// emitted once when the model decides to call one or more tools
ToolCalls(Vec<ToolCall>),
ToolCallStarted,
Completed,
Error(String),
}
// ── client ────────────────────────────────────────────────────────────────────
#[derive(Clone, Serialize, Deserialize)]
#[serde(default)]
pub struct ModelConfig {
pub api_key: String,
pub url: String,
pub model: String,
pub proxy: Option<String>,
pub extra_options: BTreeMap<String, Value>,
/// A route profile is opt-in. `None` means this route's cache contract is
/// unknown; no cache-specific field is emitted.
pub cache_profile: Option<CacheProfile>,
}
impl fmt::Debug for ModelConfig {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ModelConfig")
.field("api_key", &"<redacted>")
.field("url", &"<redacted>")
.field("model", &self.model)
.field("proxy_configured", &self.proxy.is_some())
.field("extra_option_count", &self.extra_options.len())
.field("cache_profile_configured", &self.cache_profile.is_some())
.finish()
}
}
impl Default for ModelConfig {
fn default() -> Self {
Self {
api_key: "".into(),
url: "".into(),
model: "".into(),
proxy: None,
extra_options: Default::default(),
cache_profile: None,
}
}
}
impl ModelConfig {
/// Inherited cache profiles remain bound to an unchanged provider route.
pub fn merged_over(&self, base: &ModelConfig) -> Self {
let mut extra_options = base.extra_options.clone();
extra_options.extend(self.extra_options.clone());
Self {
api_key: if self.api_key.is_empty() {
base.api_key.clone()
} else {
self.api_key.clone()
},
url: if self.url.is_empty() {
base.url.clone()
} else {
self.url.clone()
},
model: if self.model.is_empty() {
base.model.clone()
} else {
self.model.clone()
},
proxy: if self.proxy.is_some() {
self.proxy.clone()
} else {
base.proxy.clone()
},
extra_options,
cache_profile: self.cache_profile.clone().or_else(|| {
(self.url.is_empty() || self.url == base.url)
.then(|| base.cache_profile.clone())
.flatten()
}),
}
}
/// The route's extra options as the JSON object the request body takes.
fn extra_options_json(&self) -> Map<String, Value> {
self.extra_options
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default)]
pub struct ModelsConfig {
pub llm: ModelConfig,
pub embedder: ModelConfig,
pub reranker: ModelConfig,
pub embedders: Vec<ModelConfig>,
#[serde(deserialize_with = "crate::serde_utils::deserialize_usize")]
pub embed_dim: usize,
}
impl Default for ModelsConfig {
fn default() -> Self {
Self {
llm: ModelConfig {
url: "http://localhost:8001/v1".into(),
model: "auto".into(),
..Default::default()
},
embedder: ModelConfig {
url: "http://localhost:8002/v1".into(),
model: "auto".into(),
..Default::default()
},
reranker: ModelConfig {
url: "http://localhost:8003/v1".into(),
model: "auto".into(),
..Default::default()
},
embedders: Vec::new(),
embed_dim: 1024,
}
}
}
/// Line-oriented decoder for a byte stream that arrives in arbitrary chunks.
///
/// Neither boundary a chunk can land on is a real boundary: a chunk may end inside a character
/// and inside a line. Both the undecoded character and the unfinished line wait here, and this is
/// the only place stream bytes are turned into lines.
#[derive(Default)]
struct SseDecoder {
pending: Vec<u8>,
line: String,
}
impl SseDecoder {
fn feed(&mut self, chunk: &[u8]) {
self.line.push_str(&decode_chunk(&mut self.pending, chunk));
}
/// The next complete line, without its ending. `None` means the chunk ended mid-line.
fn next_line(&mut self) -> Option<String> {
let end = self.line.find('\n')?;
let line: String = self.line.drain(..=end).collect();
Some(line.trim_end_matches(['\r', '\n']).to_string())
}
/// What a stream that never ended its last line left behind.
fn unfinished(&mut self) -> String {
std::mem::take(&mut self.line)
}
}
/// Decode the completed part of a byte stream that may split a character across chunks.
///
/// A chunk boundary is not a character boundary. Bytes that can never be completed are replaced,
/// and a truncated trailing character is kept for the next chunk instead of being dropped along
/// with everything else the chunk carried.
fn decode_chunk(pending: &mut Vec<u8>, chunk: &[u8]) -> String {
pending.extend_from_slice(chunk);
if let Ok(text) = std::str::from_utf8(pending) {
let decoded = text.to_string();
pending.clear();
return decoded;
}
let error = std::str::from_utf8(pending).expect_err("checked above");
let valid = error.valid_up_to();
let invalid = error.error_len().unwrap_or(0);
let decoded = String::from_utf8_lossy(&pending[..valid + invalid]).into_owned();
pending.drain(..valid + invalid);
decoded
}
fn build_client(config: &ModelConfig) -> Result<Client> {
let accept_invalid_certs = should_accept_invalid_localhost_certs(&config.url);
let mut builder = reqwest::Client::builder()
.danger_accept_invalid_certs(accept_invalid_certs)
.user_agent("claude-code/1.0.0");
if let Some(ref proxy_url) = config.proxy {
if !proxy_url.is_empty() {
let proxy = reqwest::Proxy::all(proxy_url)?;
builder = builder.proxy(proxy);
}
}
builder.build().map_err(Into::into)
}
#[derive(Clone)]
pub struct LlmClient {
llm_client: Client,
reranker_client: Client,
embedder_clients: Vec<Client>,
embedder_index: std::sync::Arc<std::sync::atomic::AtomicUsize>,
pub config: ModelsConfig,
}
// accumulator for streaming tool call assembly
#[derive(Default)]
struct PartialCall {
id: String,
name: String,
arguments: String,
}
async fn emit_stream_terminal(
tok_tx: &mpsc::Sender<LlmEvent>,
partial_calls: HashMap<usize, PartialCall>,
final_usage: Option<Usage>,
start_time: std::time::Instant,
first_response_time: Option<std::time::Duration>,
bridge_setup_ms: u64,
) {
if !partial_calls.is_empty() {
let mut calls: Vec<_> = partial_calls.into_iter().collect();
calls.sort_by_key(|(index, _)| *index);
let calls = calls
.into_iter()
.map(|(_, call)| ToolCall {
id: call.id,
kind: "function".into(),
function: ToolCallFunction {
name: call.name,
arguments: call.arguments,
},
})
.collect();
let _ = tok_tx.send(LlmEvent::ToolCalls(calls)).await;
}
let elapsed = start_time.elapsed();
let mut usage = final_usage.unwrap_or_default();
if let Some(time_to_first_chunk) = first_response_time {
usage.prompt_processing_ms = time_to_first_chunk.as_millis() as u64;
usage.generation_ms = elapsed.saturating_sub(time_to_first_chunk).as_millis() as u64;
} else {
usage.prompt_processing_ms = elapsed.as_millis() as u64;
usage.generation_ms = 0;
}
usage.bridge_processing_ms = bridge_setup_ms;
let _ = tok_tx.send(LlmEvent::Usage(usage)).await;
let _ = tok_tx.send(LlmEvent::Completed).await;
}
impl LlmClient {
pub fn new(mut config: ModelsConfig) -> Self {
let llm_client = build_client(&config.llm).unwrap_or_else(|_| Client::new());
let reranker_client = build_client(&config.reranker).unwrap_or_else(|_| Client::new());
let mut embedder_clients = Vec::new();
if config.embedders.is_empty() {
let client = build_client(&config.embedder).unwrap_or_else(|_| Client::new());
embedder_clients.push(client);
config.embedders.push(config.embedder.clone());
} else {
for ec in &config.embedders {
let client = build_client(ec).unwrap_or_else(|_| Client::new());
embedder_clients.push(client);
}
}
Self {
llm_client,
reranker_client,
embedder_clients,
embedder_index: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
config,
}
}
fn is_auto_model(name: &str) -> bool {
let normalized = name.trim();
normalized.is_empty() || normalized.eq_ignore_ascii_case("auto")
}
fn client_for_config(&self, config: &ModelConfig) -> Client {
if config.url == self.config.llm.url {
self.llm_client.clone()
} else if config.url == self.config.reranker.url {
self.reranker_client.clone()
} else if let Some((_, c)) = self
.config
.embedders
.iter()
.zip(&self.embedder_clients)
.find(|(ec, _)| ec.url == config.url)
{
c.clone()
} else {
build_client(config).unwrap_or_else(|_| Client::new())
}
}
fn endpoint(base_url: &str, path: &str) -> String {
format!(
"{}/{}",
base_url.trim_end_matches('/'),
path.trim_start_matches("/")
)
}
fn with_auth(&self, request: RequestBuilder, config: &ModelConfig) -> RequestBuilder {
if config.api_key.is_empty() {
request
} else {
request.bearer_auth(&config.api_key)
}
}
async fn error_for_status_with_body(response: reqwest::Response) -> Result<reqwest::Response> {
let status = response.status();
if status.is_success() {
return Ok(response);
}
let url = response.url().to_string();
let body = response.text().await.unwrap_or_default();
let body = body.trim();
if body.is_empty() {
anyhow::bail!("HTTP status {status} for url ({url})");
}
let body = truncate_for_error(body, 1_000);
anyhow::bail!("HTTP status {status} for url ({url}): {body}");
}
async fn fetch_models(&self, config: &ModelConfig) -> Result<Vec<String>> {
let url = Self::endpoint(&config.url, "/models");
let client = self.client_for_config(config);
let v = self.with_auth(client.get(url), config).send().await?;
let v = Self::error_for_status_with_body(v)
.await?
.json::<Value>()
.await?;
let ids = v["data"]
.as_array()
.map(|arr| {
arr.iter()
.filter_map(|m| m["id"].as_str().map(ToString::to_string))
.collect::<Vec<_>>()
})
.unwrap_or_default();
Ok(ids)
}
fn looks_like_embedding_model(id: &str) -> bool {
let n = id.to_ascii_lowercase();
[
"embed",
"embedding",
"text-embedding",
"nomic",
"bge",
"e5",
"gte",
"mxbai",
"minilm",
"jina-emb",
]
.iter()
.any(|needle| n.contains(needle))
}
fn choose_auto_model(candidates: &[String], for_embedding: bool) -> Option<String> {
if candidates.is_empty() {
return None;
}
if for_embedding {
if let Some(id) = candidates
.iter()
.find(|id| Self::looks_like_embedding_model(id))
{
return Some(id.clone());
}
} else if let Some(id) = candidates
.iter()
.find(|id| !Self::looks_like_embedding_model(id))
{
return Some(id.clone());
}
Some(candidates[0].clone())
}
async fn resolve_model(&self, config: &ModelConfig, for_embedding: bool) -> Result<String> {
if !Self::is_auto_model(&config.model) {
return Ok(config.model.clone());
}
// Not every llama-server compatible backend implements `GET /v1/models`.
// For local/self-hosted setups, model selection is often "whatever is loaded".
// If auto-discovery fails, fall back to the configured value (usually "auto") so
// the request can still be attempted against `/v1/chat/completions`.
let models = match self.fetch_models(config).await {
Ok(models) => models,
Err(e) => {
tracing::warn!(
err = %e,
base_url = %config.url,
"auto model discovery failed; falling back to configured model id"
);
let fallback = config.model.trim();
return Ok(if fallback.is_empty() {
"auto".to_string()
} else {
fallback.to_string()
});
}
};
let selected = Self::choose_auto_model(&models, for_embedding)
.ok_or_else(|| anyhow::anyhow!("no models returned from {}/v1/models", config.url))?;
Ok(selected)
}
pub async fn stream(
&self,
messages: &[Message],
tools: &[ToolDef],
tok_tx: mpsc::Sender<LlmEvent>,
) -> Result<()> {
let body = self.stream_request_body(messages, tools).await?;
let endpoint = Self::endpoint(&self.config.llm.url, "/chat/completions");
let model = body["model"].as_str().unwrap_or("unknown");
let tool_names: Vec<&str> = tools.iter().map(|t| t.function.name.as_str()).collect();
tracing::info!(
"sending chat completions stream request to {} (model: {}, message_count: {}, tools: {:?})",
endpoint,
model,
messages.len(),
tool_names
);
let start_time = std::time::Instant::now();
let mut first_response_time: Option<std::time::Duration> = None;
let mut final_usage: Option<Usage> = None;
let res = self
.with_auth(self.llm_client.post(&endpoint), &self.config.llm)
.json(&body)
.send()
.await;
let res = match res {
Ok(r) => r,
Err(e) => {
let err_msg = format!("failed to send stream request to {endpoint}: {e:?}");
tracing::error!("{}", err_msg);
let _ = tok_tx.send(LlmEvent::Error(err_msg)).await;
return Err(e.into());
}
};
tracing::info!("stream response status: {}", res.status());
let bridge_setup_ms = res
.headers()
.get("x-bridge-setup-ms")
.and_then(|h| h.to_str().ok())
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(0);
let res = match Self::error_for_status_with_body(res).await {
Ok(r) => r,
Err(e) => {
let err_msg = format!("stream error response from {endpoint}: {e:?}");
tracing::error!("{}", err_msg);
let _ = tok_tx.send(LlmEvent::Error(err_msg)).await;
return Err(e);
}
};
let mut res = res.bytes_stream();
let mut decoder = SseDecoder::default();
let mut partial_calls: HashMap<usize, PartialCall> = HashMap::new();
let mut tool_call_started_sent = false;
let mut terminal_finish_reason: Option<String> = None;
while let Some(chunk) = res.next().await {
let chunk = match chunk {
Ok(c) => c,
Err(e) => {
let err_msg = format!("error reading stream chunk: {e:?}");
tracing::error!("{}", err_msg);
let _ = tok_tx.send(LlmEvent::Error(err_msg)).await;
return Err(e.into());
}
};
decoder.feed(&chunk);
while let Some(line) = decoder.next_line() {
tracing::debug!("stream line read: {}", line);
let Some(data) = sse_data_payload(&line) else {
continue;
};
if data == "[DONE]" {
tracing::debug!("stream line received [DONE]");
emit_stream_terminal(
&tok_tx,
partial_calls,
final_usage,
start_time,
first_response_time,
bridge_setup_ms,
)
.await;
return Ok(());
}
let v = match serde_json::from_str::<Value>(data) {
Ok(val) => val,
Err(e) => {
tracing::error!(
"failed to parse JSON from data payload: '{}', error: {:?}",
data,
e
);
continue;
}
};
// tracing::info!("got stream chunk: {v:#?}");
if v["usage"].is_object() {
let mut u: Usage = serde_json::from_value(v["usage"].clone())?;
let elapsed = start_time.elapsed();
if let Some(ttfc) = first_response_time {
u.prompt_processing_ms = ttfc.as_millis() as u64;
u.generation_ms = elapsed.saturating_sub(ttfc).as_millis() as u64;
} else {
u.prompt_processing_ms = elapsed.as_millis() as u64;
}
u.bridge_processing_ms = bridge_setup_ms;
final_usage = Some(u.clone());
let _ = tok_tx.send(LlmEvent::Usage(u)).await;
continue;
}
let Some(choices) = v["choices"].as_array() else {
continue;
};
if choices.is_empty() {
continue;
}
if let Some(fr) = choices[0].get("finish_reason").and_then(|v| v.as_str()) {
if !fr.is_empty() && fr != "null" {
terminal_finish_reason = Some(fr.to_string());
}
}
let delta = &choices[0]["delta"];
// accumulate tool call deltas
if let Some(tc_arr) = delta["tool_calls"].as_array() {
if first_response_time.is_none() {
first_response_time = Some(start_time.elapsed());
}
if !tool_call_started_sent {
let _ = tok_tx.send(LlmEvent::ToolCallStarted).await;
tool_call_started_sent = true;
}
for tc in tc_arr {
let idx = tc["index"].as_u64().unwrap_or(0) as usize;
let entry = partial_calls.entry(idx).or_default();
if let Some(id) = tc["id"].as_str() {
entry.id = id.to_string();
}
if let Some(name) = tc["function"]["name"].as_str() {
entry.name = name.to_string();
}
if let Some(args) = tc["function"]["arguments"].as_str() {
entry.arguments.push_str(args);
}
}
continue;
}
if let Some(t) = delta
.get("reasoning_content")
// this is for openrouter
.or_else(|| delta.get("reasoning"))
.and_then(|v| v.as_str())
{
if !t.is_empty() {
if first_response_time.is_none() {
first_response_time = Some(start_time.elapsed());
}
tracing::debug!("thinking token: {}", t);
if tok_tx
.send(LlmEvent::ThinkToken(t.to_string()))
.await
.is_err()
{
return Ok(());
}
}
}
if let Some(t) = delta["content"].as_str() {
if !t.is_empty() {
if first_response_time.is_none() {
first_response_time = Some(start_time.elapsed());
}
tracing::debug!("content token: {}", t);
if tok_tx.send(LlmEvent::Token(t.to_string())).await.is_err() {
return Ok(());
}
}
}
}
}
let trailing = decoder.unfinished();
if !trailing.trim().is_empty() {
let line = trailing.trim().to_string();
if let Some(data) = sse_data_payload(&line) {
if data == "[DONE]" {
emit_stream_terminal(
&tok_tx,
partial_calls,
final_usage,
start_time,
first_response_time,
bridge_setup_ms,
)
.await;
return Ok(());
}
}
}
if terminal_finish_reason.is_none() {
let err_msg = "stream terminated prematurely without [DONE]".to_string();
tracing::error!("{}", err_msg);
let _ = tok_tx.send(LlmEvent::Error(err_msg.clone())).await;
return Err(anyhow::anyhow!(err_msg));
}
emit_stream_terminal(
&tok_tx,
partial_calls,
final_usage,
start_time,
first_response_time,
bridge_setup_ms,
)
.await;
Ok(())
}
pub async fn stream_request_body(
&self,
messages: &[Message],
tools: &[ToolDef],
) -> Result<Value> {
let chat_model = self.resolve_model(&self.config.llm, false).await?;
let messages = prepare_messages_for_api(messages);
let mut body = json!({
"model": chat_model,
"messages": messages,
"stream": true,
"stream_options": { "include_usage": true },
});
body.as_object_mut()
.unwrap()
.extend(self.config.llm.extra_options_json());
self.config
.llm
.cache_profile
.clone()
.unwrap_or_default()
.apply_to_request(body.as_object_mut().unwrap());
if !tools.is_empty() {
body["tools"] = json!(tools);
}
Ok(body)
}
/// non-streaming completion, used for short internal calls with no tools.
pub async fn complete(&self, messages: &[Message]) -> Result<(String, Usage)> {
let chat_model = self.resolve_model(&self.config.llm, false).await?;
let messages = prepare_messages_for_api(messages);
let mut body = json!({
"model": chat_model,
"messages": messages,
"stream": false
});
body.as_object_mut()
.unwrap()
.extend(self.config.llm.extra_options_json());
self.config
.llm
.cache_profile
.clone()
.unwrap_or_default()
.apply_to_request(body.as_object_mut().unwrap());
let endpoint = Self::endpoint(&self.config.llm.url, "/chat/completions");
tracing::info!(
"sending chat completions non-stream request to {} (model: {}, message_count: {})",
endpoint,
chat_model,
messages.len()
);
let v = self
.with_auth(self.llm_client.post(&endpoint), &self.config.llm)
.json(&body)
.send()
.await;
let v = match v {
Ok(r) => r,
Err(e) => {
tracing::error!("failed to send non-stream request to {}: {:?}", endpoint, e);
return Err(e.into());
}
};
tracing::info!("non-stream response status: {}", v.status());
let text = Self::error_for_status_with_body(v).await?.text().await?;
if text.trim().starts_with("data:") {
Ok(complete_from_sse_text(&text))
} else {
let v: Value = serde_json::from_str(&text)?;
let content = v["choices"][0]["message"]["content"]
.as_str()
.unwrap_or("")
.to_string();
let usage: Usage = serde_json::from_value(v["usage"].clone()).unwrap_or_default();
Ok((content, usage))
}
}
fn strip_media_urls(text: &str) -> String {
rewrite_media_data_urls(text, |_| Some(String::new()))
}
async fn embed_with_config(
&self,
client: &Client,
config: &ModelConfig,
cleaned_text: &str,
) -> Result<Vec<f32>> {
let embed_model = self.resolve_model(config, true).await?;
let mut body = json!({ "model": embed_model, "input": cleaned_text });
body.as_object_mut()
.unwrap()
.extend(config.extra_options_json());
let endpoint = Self::endpoint(&config.url, "/embeddings");
tracing::info!(
"sending embeddings request to {} (model: {}, input_chars: {})",
endpoint,
embed_model,
cleaned_text.chars().count()
);
let v = self
.with_auth(client.post(&endpoint), config)
.json(&body)
.send()
.await;
let v = match v {
Ok(r) => r,
Err(e) => {
tracing::error!("failed to send embeddings request to {}: {:?}", endpoint, e);
return Err(e.into());
}
};
tracing::info!("embeddings response status: {}", v.status());
let v = Self::error_for_status_with_body(v)
.await?
.json::<Value>()
.await?;
let embedding = Self::parse_single_embedding_response(&v)?;
if embedding.len() != self.config.embed_dim {
return Err(anyhow::anyhow!(
"embedding dimension mismatch: got {}, expected {} (model: {})",
embedding.len(),
self.config.embed_dim,
embed_model
));
}
Ok(embedding)
}
fn parse_embedding_item(item: &Value) -> Result<Vec<f32>> {
let raw_embedding = item
.get("embedding")
.ok_or_else(|| anyhow::anyhow!("no embedding in response"))?;
let values = raw_embedding
.as_array()
.ok_or_else(|| anyhow::anyhow!("embedding was not an array"))?;
let values = if values.len() == 1 && values[0].is_array() {
values[0]
.as_array()
.ok_or_else(|| anyhow::anyhow!("embedding was not an array"))?
} else {
values
};
values
.iter()
.map(|x| {
x.as_f64()
.map(|value| value as f32)
.ok_or_else(|| anyhow::anyhow!("embedding value was not a number"))
})
.collect()
}
fn embedding_items(v: &Value) -> Result<&Vec<Value>> {
if let Some(data) = v.get("data").and_then(Value::as_array) {
return Ok(data);
}
v.as_array()
.ok_or_else(|| anyhow::anyhow!("no embedding data in response"))
}
fn parse_embedding_batch_response(v: &Value, expected_count: usize) -> Result<Vec<Vec<f32>>> {
let items = Self::embedding_items(v)?;
if items.len() != expected_count {
return Err(anyhow::anyhow!(
"batch embedding response count mismatch: got {}, expected {}",
items.len(),
expected_count
));
}
let mut parsed = items
.iter()
.enumerate()
.map(|(pos, item)| {
let index = item
.get("index")
.and_then(Value::as_u64)
.map(|index| index as usize)
.unwrap_or(pos);
Ok((index, Self::parse_embedding_item(item)?))
})
.collect::<Result<Vec<_>>>()?;
parsed.sort_by_key(|(index, _)| *index);
Ok(parsed.into_iter().map(|(_, embedding)| embedding).collect())
}
fn parse_single_embedding_response(v: &Value) -> Result<Vec<f32>> {
Self::parse_embedding_batch_response(v, 1)?
.into_iter()
.next()
.ok_or_else(|| anyhow::anyhow!("no embedding in response"))
}
pub async fn embed(&self, text: &str) -> Result<Vec<f32>> {
let cleaned_text = Self::strip_media_urls(text);
let num_embedders = self.config.embedders.len();
if num_embedders == 0 {
return Err(anyhow::anyhow!("no embedders configured"));
}
let start_idx = self
.embedder_index
.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
% num_embedders;
let mut last_err = None;
for i in 0..num_embedders {
let idx = (start_idx + i) % num_embedders;
let config = &self.config.embedders[idx];
let client = &self.embedder_clients[idx];
match self.embed_with_config(client, config, &cleaned_text).await {
Ok(embedding) => return Ok(embedding),
Err(e) => {
tracing::warn!(
index = idx,
url = %config.url,
err = %e,
"embedding attempt failed; trying next embedder if available"
);
last_err = Some(e);
}
}
}
Err(last_err.unwrap_or_else(|| anyhow::anyhow!("no embedders configured or available")))
}
async fn embed_chunk_with_config(
&self,
client: &Client,
config: &ModelConfig,
chunk_texts: &[String],
) -> Result<Vec<Vec<f32>>> {
let embed_model = self.resolve_model(config, true).await?;
let endpoint = Self::endpoint(&config.url, "/embeddings");
let body = json!({
"model": embed_model,
"input": chunk_texts,
});
tracing::info!(
"sending batch embeddings request chunk to {} (model: {}, count: {})",
endpoint,
embed_model,
chunk_texts.len()
);
let response = self
.with_auth(client.post(&endpoint), config)
.json(&body)
.send()
.await?;
let response = Self::error_for_status_with_body(response).await?;
let parsed = response.json::<Value>().await?;
let embeddings = Self::parse_embedding_batch_response(&parsed, chunk_texts.len())?;
let mut results = Vec::with_capacity(chunk_texts.len());
for embedding in embeddings {
if embedding.len() != self.config.embed_dim {
return Err(anyhow::anyhow!(
"embedding dimension mismatch: got {}, expected {} (model: {})",
embedding.len(),
self.config.embed_dim,
embed_model
));
}
results.push(embedding);
}
Ok(results)
}
async fn embed_chunk_with_failover(
&self,
chunk_texts: Vec<String>,
chunk_indices: Vec<usize>,
) -> Result<Vec<(usize, Vec<f32>)>> {
let num_embedders = self.config.embedders.len();
if num_embedders == 0 {
return Err(anyhow::anyhow!("no embedders configured"));
}
let start_idx = self
.embedder_index
.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
% num_embedders;
let mut last_err = None;
for i in 0..num_embedders {
let idx = (start_idx + i) % num_embedders;
let config = &self.config.embedders[idx];
let client = &self.embedder_clients[idx];
match self
.embed_chunk_with_config(client, config, &chunk_texts)
.await
{
Ok(embeddings) => {
let results = chunk_indices.iter().copied().zip(embeddings).collect();
return Ok(results);
}
Err(e) => {
tracing::warn!(
index = idx,
url = %config.url,
err = %e,
"chunk embedding attempt failed; trying next embedder if available"
);
last_err = Some(e);
}
}
}
Err(last_err.unwrap_or_else(|| anyhow::anyhow!("all embedders failed for chunk")))
}
pub async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
if texts.is_empty() {
return Ok(Vec::new());
}
let cleaned_texts: Vec<String> = texts.iter().map(|t| Self::strip_media_urls(t)).collect();
// Chunking parameters
const MAX_TEXTS_PER_CHUNK: usize = 128;
const MAX_TOKENS_PER_CHUNK: usize = 16000;
let mut current_chunk_texts = Vec::new();
let mut current_chunk_indices = Vec::new();
let mut current_chunk_tokens = 0;
let mut chunks = Vec::new();
for (idx, text) in cleaned_texts.into_iter().enumerate() {
let estimated_tokens = (text.len() / 4).max(1);
if current_chunk_texts.len() >= MAX_TEXTS_PER_CHUNK
|| current_chunk_tokens + estimated_tokens > MAX_TOKENS_PER_CHUNK
{
chunks.push((
std::mem::take(&mut current_chunk_texts),
std::mem::take(&mut current_chunk_indices),
));
current_chunk_tokens = 0;
}
current_chunk_tokens += estimated_tokens;
current_chunk_texts.push(text);
current_chunk_indices.push(idx);
}
if !current_chunk_texts.is_empty() {
chunks.push((current_chunk_texts, current_chunk_indices));
}
let mut futures = Vec::new();
for (chunk_texts, chunk_indices) in chunks {
futures.push(self.embed_chunk_with_failover(chunk_texts, chunk_indices));
}
let results_by_chunk = futures::future::try_join_all(futures).await?;
let mut final_results = vec![Vec::new(); texts.len()];
for chunk_res in results_by_chunk {
for (original_idx, embedding) in chunk_res {
final_results[original_idx] = embedding;
}
}
Ok(final_results)
}
}
fn sse_data_payload(line: &str) -> Option<&str> {
let line = line.trim();
// Standard SSE uses "data: <payload>" but some providers omit the space.
line.strip_prefix("data: ")
.or_else(|| line.strip_prefix("data:"))
.map(str::trim)
}
fn complete_from_sse_text(text: &str) -> (String, Usage) {
let mut content = String::new();
let mut usage = Usage::default();
for data in text.lines().filter_map(sse_data_payload) {
if data == "[DONE]" {
continue;
}
let Ok(val) = serde_json::from_str::<Value>(data) else {
continue;
};
if let Some(choices) = val["choices"].as_array() {
if let Some(choice) = choices.first() {
if let Some(delta) = choice.get("delta") {
if let Some(chunk) = delta["content"].as_str() {
content.push_str(chunk);
}
}
if let Some(message) = choice.get("message") {
if let Some(text_content) = message["content"].as_str() {
content.push_str(text_content);
}
}
}
}
if let Some(value) = val.get("usage") {
if let Ok(parsed) = serde_json::from_value::<Usage>(value.clone()) {
usage = parsed;
}
}
}
(content, usage)
}
struct MediaDataUrl<'a> {
start: usize,
end: usize,
mime: &'a str,
approx_size: usize,
}
struct MediaScan<'a> {
data_start: usize,
media: Option<MediaDataUrl<'a>>,
}
fn next_media_scan(text: &str, current_idx: usize) -> Option<MediaScan<'_>> {
let data_start = current_idx + text[current_idx..].find("data:")?;
Some(MediaScan {
data_start,
media: parse_media_data_url_at(text, data_start),
})
}
fn parse_media_data_url_at(text: &str, start: usize) -> Option<MediaDataUrl<'_>> {
let comma_offset = text[start..].find(";base64,")?;
if comma_offset >= 100 {
return None;
}
let mime = &text[start + "data:".len()..start + comma_offset];
if mime.contains("data:") {
return None;
}
let base64_start = start + comma_offset + ";base64,".len();
let mut base64_end = base64_start;
for ch in text[base64_start..].chars() {
if ch.is_ascii_alphanumeric() || ch == '+' || ch == '/' || ch == '=' {
base64_end += ch.len_utf8();
} else {
break;
}
}
if base64_end == base64_start {
return None;
}
let base64_len = base64_end - base64_start;
Some(MediaDataUrl {
start,
end: base64_end,
mime,
approx_size: (base64_len * 3) / 4,
})
}
fn rewrite_media_data_urls(
text: &str,
mut replacement: impl FnMut(&MediaDataUrl<'_>) -> Option<String>,
) -> String {
let mut result = String::new();
let mut current_idx = 0;
while let Some(scan) = next_media_scan(text, current_idx) {
if let Some(media) = scan.media {
if let Some(value) = replacement(&media) {
result.push_str(&text[current_idx..media.start]);
result.push_str(&value);
current_idx = media.end;
continue;
}
}
result.push_str(&text[current_idx..scan.data_start + "data:".len()]);
current_idx = scan.data_start + "data:".len();
}
result.push_str(&text[current_idx..]);
result
}
fn replace_media_urls_with_placeholders(text: &str, limit_bytes: usize) -> String {
rewrite_media_data_urls(text, |media| {
if media.approx_size > limit_bytes {
Some(format!(
"[Media attachment: {}, approximate size: {} bytes]",
media.mime, media.approx_size
))
} else {
None
}
})
}
fn prepare_messages_for_api(messages: &[Message]) -> Vec<Message> {
let mut prepared = Vec::new();
for mut message in messages.iter().cloned() {
if message.normalize_for_api() {
if let Some(content) = message.content.as_mut() {
*content = replace_media_urls_with_placeholders(content, 10 * 1024 * 1024);
}
prepared.push(message);
}
}
prepared
}
fn truncate_for_error(value: &str, max_chars: usize) -> String {
let mut chars = value.chars();
let truncated: String = chars.by_ref().take(max_chars).collect();
if chars.next().is_some() {
format!("{truncated}...")
} else {
truncated
}
}
fn should_accept_invalid_localhost_certs(url: &str) -> bool {
let Ok(parsed) = reqwest::Url::parse(url) else {
return false;
};
if parsed.scheme() != "https" {
return false;
}
matches!(
parsed.host_str(),
Some("localhost") | Some("127.0.0.1") | Some("::1")
)
}
#[cfg(test)]
mod usage_shape_tests {
use super::Usage;
/// DeepSeek: the hit/miss pair sits next to `prompt_tokens`, not under a details object.
#[test]
fn a_route_that_reports_hit_and_miss_tokens_is_measured() {
let usage: Usage = serde_json::from_value(serde_json::json!({
"prompt_tokens": 35072,
"completion_tokens": 156,
"total_tokens": 35228,
"prompt_cache_hit_tokens": 35072,
"prompt_cache_miss_tokens": 0,
}))
.expect("usage");
assert_eq!(usage.cached_prompt_tokens, Some(35072));
assert_eq!(usage.uncached_prompt_tokens, Some(0));
assert_eq!(
usage.cache_write_tokens, None,
"this route never reported a write"
);
assert!(usage.tokens_reported);
}
/// OpenAI on a model with cache writes: the details object carries both halves.
#[test]
fn a_route_that_reports_writes_and_reads_is_measured_on_both_sides() {
let usage: Usage = serde_json::from_value(serde_json::json!({
"prompt_tokens": 15000,
"completion_tokens": 12,
"total_tokens": 15012,
"prompt_tokens_details": { "cached_tokens": 12000, "cache_write_tokens": 3000 },
}))
.expect("usage");
assert_eq!(usage.cached_prompt_tokens, Some(12000));
assert_eq!(usage.cache_write_tokens, Some(3000));
assert_eq!(usage.uncached_prompt_tokens, None);
}
/// A route that says nothing must not read as a route that said zero.
#[test]
fn a_route_that_says_nothing_leaves_the_measurements_absent() {
let usage: Usage = serde_json::from_value(serde_json::json!({
"prompt_tokens": 900,
"completion_tokens": 4,
"total_tokens": 904,
}))
.expect("usage");
assert_eq!(usage.cached_prompt_tokens, None);
assert_eq!(usage.cache_write_tokens, None);
assert_eq!(usage.uncached_prompt_tokens, None);
}
}
#[cfg(test)]
mod stream_chunk_tests {
use super::SseDecoder;
/// A chunk boundary lands wherever the network put it: inside a character, inside a line, or
/// both. Every split of the same stream has to produce the same lines.
#[test]
fn every_chunk_split_produces_the_same_lines() {
let text = "data: 先に 日本語 🙂\ndata: second line\n";
let expected: Vec<String> = text.lines().map(str::to_string).collect();
let bytes = text.as_bytes();
for size in 1..bytes.len() {
let mut decoder = SseDecoder::default();
let mut lines = Vec::new();
for chunk in bytes.chunks(size) {
decoder.feed(chunk);
while let Some(line) = decoder.next_line() {
lines.push(line);
}
}
let trailing = decoder.unfinished();
if !trailing.trim().is_empty() {
lines.push(trailing.trim().to_string());
}
assert_eq!(lines, expected, "chunk size {size}");
}
}
/// Bytes that can never be completed are replaced, and the rest of the stream is kept.
#[test]
fn invalid_bytes_are_replaced_without_eating_the_stream() {
let mut decoder = SseDecoder::default();
decoder.feed(&[0x41, 0xff, 0x42]);
decoder.feed(b" tail\n");
assert_eq!(decoder.next_line().as_deref(), Some("A\u{fffd}B tail"));
}
}