Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
180 changes: 148 additions & 32 deletions src-tauri/src/proxy/providers/streaming.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ struct DeltaFunction {
arguments: Option<String>,
}

#[derive(Debug, Deserialize)]
#[derive(Debug, Clone, Deserialize)]
struct Usage {
#[serde(default)]
prompt_tokens: u32,
Expand All @@ -67,7 +67,7 @@ struct Usage {
cache_creation_input_tokens: Option<u32>,
}

#[derive(Debug, Deserialize)]
#[derive(Debug, Clone, Deserialize)]
struct PromptTokensDetails {
#[serde(default)]
cached_tokens: u32,
Expand Down Expand Up @@ -98,6 +98,11 @@ pub fn create_anthropic_sse_stream(
let mut open_tool_block_indices: HashSet<u32> = HashSet::new();
let mut legacy_function_name: Option<String> = None;
let mut legacy_function_block_index: Option<u32> = None;
// Final usage and stop reason are deferred: OpenAI-compatible upstreams emit
// them in separate trailing chunks (the usage chunk has empty `choices`), so
// we cache them here and flush a single message_delta at [DONE]/stream end.
let mut latest_usage: Option<Usage> = None;
let mut pending_stop_reason: Option<String> = None;

tokio::pin!(stream);

Expand All @@ -120,6 +125,17 @@ pub fn create_anthropic_sse_stream(
};

if data.trim() == "[DONE]" {
if has_sent_message_start && pending_stop_reason.is_some() {
let event = build_message_delta_event(
pending_stop_reason.as_deref(),
latest_usage.as_ref(),
);
let sse_data = format!(
"event: message_delta\ndata: {}\n\n",
serde_json::to_string(&event).unwrap_or_default()
);
yield Ok(Bytes::from(sse_data));
}
let event = json!({"type": "message_stop"});
let sse_data = format!(
"event: message_stop\ndata: {}\n\n",
Expand All @@ -141,6 +157,15 @@ pub fn create_anthropic_sse_stream(
current_model = Some(chunk.model.clone());
}

// Cache usage from any chunk (incl. the trailing
// `choices: []` usage chunk that include_usage produces).
// Must happen before the choices.first() guard below, which
// otherwise drops that chunk and zeroes token/cost/cache
// accounting. Mirrors streaming_codex_chat's latest_usage.
if let Some(usage) = &chunk.usage {
latest_usage = Some(usage.clone());
}

let Some(choice) = chunk.choices.first() else {
continue;
};
Expand All @@ -166,6 +191,15 @@ pub fn create_anthropic_sse_stream(
"id": message_id.clone().unwrap_or_default(),
"type": "message",
"role": "assistant",
// Anthropic SSE spec requires these fields in the
// message snapshot; the official SDK's stream
// accumulator does `snapshot.content.push(...)` and
// crashes ("Cannot read properties of undefined")
// when `content` is missing, forcing clients into
// non-streaming fallback (or silent stream loss).
"content": [],
"stop_reason": serde_json::Value::Null,
"stop_sequence": serde_json::Value::Null,
"model": current_model.clone().unwrap_or_default(),
"usage": start_usage
}
Expand Down Expand Up @@ -588,36 +622,13 @@ pub fn create_anthropic_sse_stream(
open_tool_block_indices.clear();
}

let usage_json = if let Some(usage) = chunk.usage.as_ref() {
let mut usage_json = json!({
"input_tokens": usage.prompt_tokens,
"output_tokens": usage.completion_tokens
});
if let Some(cached) = extract_cache_read_tokens(usage) {
usage_json["cache_read_input_tokens"] = json!(cached);
}
if let Some(created) = usage.cache_creation_input_tokens {
usage_json["cache_creation_input_tokens"] = json!(created);
}
usage_json
} else {
json!({
"output_tokens": 0
})
};
let event = json!({
"type": "message_delta",
"delta": {
"stop_reason": map_stop_reason(Some(finish_reason)),
"stop_sequence": null
},
"usage": usage_json
});
let sse_data = format!(
"event: message_delta\ndata: {}\n\n",
serde_json::to_string(&event).unwrap_or_default()
);
yield Ok(Bytes::from(sse_data));
// Defer the message_delta: the final usage often arrives
// in a *later* chunk than finish_reason (OpenAI-compatible
// streams send `choices:[{finish_reason}]` then a separate
// `choices:[]` usage chunk). Emitting now would lock in
// output_tokens=0; instead record the stop reason and flush
// one message_delta with the cached usage at stream end.
pending_stop_reason = Some(finish_reason.clone());
}
}
}
Expand All @@ -641,10 +652,53 @@ pub fn create_anthropic_sse_stream(
}
}

// Flush the deferred message_delta if the upstream closed without a [DONE]
// marker, so the final usage still reaches the client.
if has_sent_message_start && pending_stop_reason.is_some() {
let event = build_message_delta_event(
pending_stop_reason.as_deref(),
latest_usage.as_ref(),
);
let sse_data = format!(
"event: message_delta\ndata: {}\n\n",
serde_json::to_string(&event).unwrap_or_default()
);
yield Ok(Bytes::from(sse_data));
}

stream_completion.record_success();
}
}

fn build_message_delta_event(
stop_reason: Option<&str>,
usage: Option<&Usage>,
) -> serde_json::Value {
let usage_json = if let Some(usage) = usage {
let mut usage_json = json!({
"input_tokens": usage.prompt_tokens,
"output_tokens": usage.completion_tokens
});
if let Some(cached) = extract_cache_read_tokens(usage) {
usage_json["cache_read_input_tokens"] = json!(cached);
}
if let Some(created) = usage.cache_creation_input_tokens {
usage_json["cache_creation_input_tokens"] = json!(created);
}
usage_json
} else {
json!({ "output_tokens": 0 })
};
json!({
"type": "message_delta",
"delta": {
"stop_reason": map_stop_reason(stop_reason),
"stop_sequence": null
},
"usage": usage_json
})
}

fn map_stop_reason(finish_reason: Option<&str>) -> Option<String> {
finish_reason.map(|reason| {
match reason {
Expand Down Expand Up @@ -819,6 +873,68 @@ mod tests {
);
}

#[tokio::test]
async fn streaming_usage_from_trailing_empty_choices_chunk_is_mapped() {
// OpenAI-compatible upstreams (LiteLLM, kimi, MiniMax, ...) send the final
// usage in a separate `choices: []` chunk *after* the finish_reason chunk
// when stream_options.include_usage=true. The choices.first() guard used to
// drop it, leaving message_delta usage at 0, zeroing token/cost accounting.
let input = concat!(
"data: {\"id\":\"chatcmpl_1\",\"model\":\"gpt-5.5\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
"data: {\"id\":\"chatcmpl_1\",\"model\":\"gpt-5.5\",\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
"data: {\"id\":\"chatcmpl_1\",\"model\":\"gpt-5.5\",\"choices\":[],\"usage\":{\"prompt_tokens\":12,\"completion_tokens\":9,\"total_tokens\":21}}\n\n",
"data: [DONE]\n\n"
);

let events = collect_events(input).await;
let message_delta = events
.iter()
.find(|event| event["type"] == "message_delta")
.expect("message_delta event");

assert_eq!(message_delta["delta"]["stop_reason"], "end_turn");
assert_eq!(message_delta["usage"]["input_tokens"], 12);
assert_eq!(message_delta["usage"]["output_tokens"], 9);

// Exactly one message_delta is emitted, and it precedes message_stop.
let delta_count = events
.iter()
.filter(|event| event["type"] == "message_delta")
.count();
assert_eq!(delta_count, 1);
let names: Vec<&str> = events
.iter()
.filter_map(|event| event["type"].as_str())
.collect();
let delta_pos = names.iter().position(|n| *n == "message_delta");
let stop_pos = names.iter().position(|n| *n == "message_stop");
assert!(
delta_pos < stop_pos,
"message_delta must precede message_stop"
);
}

#[tokio::test]
async fn streaming_message_start_is_spec_complete() {
// message_start must carry content/stop_reason/stop_sequence or the
// official Anthropic SDK stream accumulator crashes on snapshot.content.
let input = concat!(
"data: {\"id\":\"chatcmpl_1\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
"data: {\"id\":\"chatcmpl_1\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1}}\n\n",
"data: [DONE]\n\n"
);

let events = collect_events(input).await;
let message = &events
.iter()
.find(|event| event["type"] == "message_start")
.expect("message_start event")["message"];

assert_eq!(message["content"], serde_json::json!([]));
assert_eq!(message["stop_reason"], serde_json::Value::Null);
assert_eq!(message["stop_sequence"], serde_json::Value::Null);
}

#[tokio::test]
async fn streaming_accepts_deepseek_reasoning_content_alias() {
let input = concat!(
Expand Down
33 changes: 33 additions & 0 deletions src-tauri/src/proxy/providers/streaming_gemini.rs
Original file line number Diff line number Diff line change
Expand Up @@ -316,6 +316,11 @@ pub fn create_anthropic_sse_stream_from_gemini<E: std::error::Error + Send + 'st
"id": message_id.clone().unwrap_or_default(),
"type": "message",
"role": "assistant",
// Spec-required fields; missing `content` crashes the
// official Anthropic SDK stream accumulator.
"content": [],
"stop_reason": serde_json::Value::Null,
"stop_sequence": serde_json::Value::Null,
"model": current_model.clone().unwrap_or_default(),
"usage": build_anthropic_usage(chunk_json.get("usageMetadata"))
}
Expand Down Expand Up @@ -423,6 +428,11 @@ pub fn create_anthropic_sse_stream_from_gemini<E: std::error::Error + Send + 'st
"id": message_id.clone().unwrap_or_default(),
"type": "message",
"role": "assistant",
// Spec-required fields; missing `content` crashes the
// official Anthropic SDK stream accumulator.
"content": [],
"stop_reason": serde_json::Value::Null,
"stop_sequence": serde_json::Value::Null,
"model": current_model.clone().unwrap_or_default(),
"usage": build_anthropic_usage(latest_usage.as_ref())
}
Expand Down Expand Up @@ -657,6 +667,29 @@ mod tests {
assert!(output.contains("event: message_stop"));
}

#[test]
fn message_start_is_spec_complete() {
// message_start must carry content/stop_reason/stop_sequence or the
// official Anthropic SDK stream accumulator crashes on snapshot.content.
let output = collect_stream_output(vec![
"data: {\"responseId\":\"resp_1\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"finishReason\":\"STOP\",\"content\":{\"parts\":[{\"text\":\"Hi\"}]}}],\"usageMetadata\":{\"promptTokenCount\":4,\"totalTokenCount\":6}}\n\n",
]);

let data = output
.split("\n\n")
.find_map(|block| {
let line = block.lines().find(|l| l.starts_with("data: "))?;
let value: Value = serde_json::from_str(line.trim_start_matches("data: ")).ok()?;
(value["type"] == "message_start").then_some(value)
})
.expect("message_start event");
let message = &data["message"];

assert_eq!(message["content"], serde_json::json!([]));
assert_eq!(message["stop_reason"], Value::Null);
assert_eq!(message["stop_sequence"], Value::Null);
}

#[test]
fn records_stream_completion_success() {
let stream = futures::stream::iter([Ok::<Bytes, std::io::Error>(Bytes::from_static(
Expand Down
43 changes: 43 additions & 0 deletions src-tauri/src/proxy/providers/streaming_responses.rs
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,11 @@ pub fn create_anthropic_sse_stream_from_responses(
"id": message_id.clone().unwrap_or_default(),
"type": "message",
"role": "assistant",
// Spec-required fields; missing `content` crashes the
// official Anthropic SDK stream accumulator.
"content": [],
"stop_reason": serde_json::Value::Null,
"stop_sequence": serde_json::Value::Null,
"model": current_model.clone().unwrap_or_default(),
"usage": build_anthropic_usage_from_responses(response_obj.get("usage"))
}
Expand All @@ -185,6 +190,11 @@ pub fn create_anthropic_sse_stream_from_responses(
"id": message_id.clone().unwrap_or_default(),
"type": "message",
"role": "assistant",
// Spec-required fields; missing `content` crashes the
// official Anthropic SDK stream accumulator.
"content": [],
"stop_reason": serde_json::Value::Null,
"stop_sequence": serde_json::Value::Null,
"model": current_model.clone().unwrap_or_default(),
"usage": { "input_tokens": 0, "output_tokens": 0 }
}
Expand Down Expand Up @@ -301,6 +311,11 @@ pub fn create_anthropic_sse_stream_from_responses(
"id": message_id.clone().unwrap_or_default(),
"type": "message",
"role": "assistant",
// Spec-required fields; missing `content` crashes the
// official Anthropic SDK stream accumulator.
"content": [],
"stop_reason": serde_json::Value::Null,
"stop_sequence": serde_json::Value::Null,
"model": current_model.clone().unwrap_or_default(),
"usage": { "input_tokens": 0, "output_tokens": 0 }
}
Expand Down Expand Up @@ -748,6 +763,34 @@ data: {\"response\":{\"status\":\"completed\"}}\r\n\
assert!(merged.contains("\"stop_reason\":\"tool_use\""));
}

#[tokio::test]
async fn tool_first_message_start_is_spec_complete() {
// A response that opens with a function_call (no preceding text) emits
// message_start from the tool path; it must still carry the spec-required
// fields or the official Anthropic SDK stream accumulator crashes.
let input = concat!(
"event: response.created\n",
"data: {\"response\":{\"id\":\"resp_tool\",\"model\":\"gpt-5.4\"}}\n\n",
"event: response.output_item.added\n",
"data: {\"item\":{\"id\":\"fc_1\",\"type\":\"function_call\",\"call_id\":\"call_1\",\"name\":\"get_weather\"}}\n\n",
"event: response.function_call_arguments.done\n",
"data: {\"item_id\":\"fc_1\",\"arguments\":\"{}\"}\n\n",
"event: response.completed\n",
"data: {\"response\":{\"status\":\"completed\"}}\n\n"
);

let (merged, _) = collect_stream(vec![Bytes::from(input)]).await;
let events = parse_anthropic_events(&merged);
let message = &events
.iter()
.find(|event| event["type"] == "message_start")
.expect("message_start event")["message"];

assert_eq!(message["content"], serde_json::json!([]));
assert_eq!(message["stop_reason"], Value::Null);
assert_eq!(message["stop_sequence"], Value::Null);
}

#[tokio::test]
async fn non_read_tool_empty_delta_still_uses_done_arguments() {
let input = concat!(
Expand Down
Loading