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
6 changes: 5 additions & 1 deletion desktop-tauri/src-tauri/server/src/server/api/protocol.rs
Original file line number Diff line number Diff line change
Expand Up @@ -419,6 +419,10 @@ pub async fn messages_endpoint(
telemetry.set_cancel_token(token);
}
let mut guard = DisconnectGuard::new(state.request_stats(), telemetry_id.clone());
// `message_start` 必须立刻发、而上游 usage 要到流末尾才到;这里先按
// 请求体估一个输入 token 数占位(真实值在 `message_delta` 覆盖)。
// 纯内存扫一遍消息 + 工具定义,不增加任何网络等待。
let prompt_estimate = anthropic::estimate_chat_tokens(&chat_body);
let outcome = forward_chat(
&state,
chat_body,
Expand Down Expand Up @@ -455,7 +459,7 @@ pub async fn messages_endpoint(
let model = requested_model.clone();
let context = RecordContext { status: i64::from(status.as_u16()), ..context };
if stream {
let mut machine = anthropic::AnthropicStream::new(&model);
let mut machine = anthropic::AnthropicStream::with_input_tokens(&model, prompt_estimate);
let transformed = pipeline::transformed_stream(
source,
context,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -676,6 +676,35 @@ pub(super) fn parse_json_object(value: &Value) -> Value {
}

/// Chat 的 usage → Anthropic 的 usage
/// 输入 token 的粗估(CJK 一字一 token,其余约 4 字符一 token)。
///
/// 只用于 `message_start` 的占位:上游真实用量要到流末尾才出现,而
/// `message_start` 必须在首帧就发。真实值随后在 `message_delta` 覆盖。
pub fn estimate_chat_tokens(chat: &Value) -> i64 {
let payload = json!({
"messages": chat.get("messages").cloned().unwrap_or_else(|| json!([])),
"tools": chat.get("tools").cloned().unwrap_or_else(|| json!([])),
});
let text = payload.to_string();
let (mut cjk, mut total) = (0i64, 0i64);
for ch in text.chars() {
total += 1;
if is_cjk_char(ch) {
cjk += 1;
}
}
cjk + ((total - cjk) + 3) / 4
}

fn is_cjk_char(ch: char) -> bool {
matches!(ch as u32,
0x1100..=0x11ff
| 0x2e80..=0x9fff
| 0xac00..=0xd7af
| 0xf900..=0xfaff
| 0xff00..=0xffef)
}

pub fn usage_to_anthropic(usage: Option<&Value>) -> Value {
let Some(usage) = usage.filter(|value| value.is_object()) else {
return json!({ "input_tokens": 0, "output_tokens": 0 });
Expand Down Expand Up @@ -731,6 +760,12 @@ pub struct AnthropicStream {

impl AnthropicStream {
pub fn new(model: &str) -> Self {
Self::with_input_tokens(model, 0)
}

/// `input_tokens` 用于 `message_start` 的占位(上游真实用量要到流末尾
/// 才出现);`message_delta` 会用真实值覆盖。
pub fn with_input_tokens(model: &str, input_tokens: i64) -> Self {
Self {
buffer: SseLineBuffer::new(),
message_id: random_id("msg"),
Expand All @@ -748,7 +783,7 @@ impl AnthropicStream {
has_tool: false,
stop_reason: None,
usage: None,
input_tokens: 0,
input_tokens,
}
}

Expand Down Expand Up @@ -789,25 +824,22 @@ impl AnthropicStream {
.stop_reason
.clone()
.unwrap_or_else(|| if self.has_tool { "tool_use".to_string() } else { "end_turn".to_string() });
// 真实 Anthropic 流把 `input_tokens` / `cache_read_input_tokens` 放在
// 收尾的 `message_delta`(开头那帧只报 0 占位,见
// `anthropic_outbound::ChatFromAnthropicStream::absorb_usage`)。
// 只回传 output_tokens 会让客户端(含 Claude Code 的 context-hud)
// 认为本次请求没有输入,上下文用量永远是 0。
out.push(self.event(
"message_delta",
json!({
"delta": { "stop_reason": stop_reason, "stop_sequence": Value::Null },
"usage": { "output_tokens": self.output_tokens() },
"usage": usage_to_anthropic(self.usage.as_ref()),
}),
));
out.push(self.event("message_stop", json!({})));
out
}

fn output_tokens(&self) -> i64 {
self.usage
.as_ref()
.and_then(|usage| usage.get("completion_tokens"))
.and_then(Value::as_i64)
.unwrap_or(0)
}

fn consume(&mut self, chunk: &Value) -> Vec<bytes::Bytes> {
if self.finished {
return Vec::new();
Expand All @@ -833,7 +865,8 @@ impl AnthropicStream {
}
}
if let Some(usage) = chunk.get("usage").filter(|value| value.is_object()) {
// Anthropic 的 message_delta 只报 output_tokens,input 在 message_start
// 上游的 usage 只在流末尾出现;`input_tokens` 留到收尾的
// `message_delta` 一起下发(见 `finish`)。
self.input_tokens = usage.get("prompt_tokens").and_then(Value::as_i64).unwrap_or(0);
self.usage = Some(usage.clone());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -447,6 +447,11 @@ pub fn build_upstream_body(
"version": "1.0.0",
"type": "agent",
"stage": "start",
// 官方 agent 循环的 business 固定带 `model_key`(官方运行时 `d8()`
// 的第 5 个参数就是当前模型键)。上游后台的模型归属/计费优先读它,
// 缺了会把请求记到账号默认档 Auto —— 即使 model_config.key 与
// X-Model-Key 都是正确的。
"model_key": upstream_key,
"id": super::cosy::random_uuid().unwrap_or_default(),
"name": last_user_text.chars().take(30).collect::<String>(),
"begin_at": crate::server::logging::now_ms(),
Expand Down Expand Up @@ -555,3 +560,36 @@ pub fn truthy(value: &Value) -> bool {
Value::Array(_) | Value::Object(_) => true,
}
}

#[cfg(test)]
mod tests {
use super::*;

/// 上游后台按 `business.model_key` 归属模型;缺了会把请求记成账号默认档
/// Auto(即使 `model_config.key` / `X-Model-Key` 正确)。
#[test]
fn business_carries_requested_model_key() {
let messages = vec![json!({ "role": "user", "content": "hi" })];
let body = build_upstream_body(
"kmodel_latest",
&json!({ "key": "kmodel_latest", "source": "system" }),
&messages,
None,
Some(16),
&ThinkingChoice {
enable: None,
effort: None,
},
"uid",
None,
);
assert_eq!(
body.pointer("/business/model_key").and_then(Value::as_str),
Some("kmodel_latest")
);
assert_eq!(
body.pointer("/model_config/key").and_then(Value::as_str),
Some("kmodel_latest")
);
}
}