Dark-Alex-17/coyote
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Support token caching to reduce LLM costs in Anthropic models (not typically handled automatically by APIs)

Browse Source
This commit is contained in:
Dark-Alex-17 committed 2026-09-14 13:56:13 -06:00
1 parent 9a6ffcb8ca
commit 5d59adadbc
10 files changed
+984 -78
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/client/bedrock.rs
+361 -39
View File
@@ -1,3 +1,4 @@
use std::collections::HashSet;
use std::mem;
use super::*;
@@ -24,6 +25,7 @@ pub struct BedrockConfig {
pub session_token: Option<String>,
#[serde(default)]
pub models: Vec<ModelData>,
pub prompt_cache: Option<bool>,
pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>,
}
@@ -59,7 +61,11 @@ impl BedrockClient {
format!("/model/{model_name}/converse")
};
let body = build_chat_completions_body(data, &self.model)?;
let body = build_chat_completions_body(
data,
&self.model,
self.config.prompt_cache.unwrap_or(true),
)?;
let mut request_data = RequestData::new("", body);
self.patch_request_data(&mut request_data);
@@ -284,6 +290,12 @@ async fn chat_completions_streaming(
function_arguments.clear();
}
}
"metadata" => {
if let Some(usage) = bedrock_parse_usage(&data["usage"]) {
debug!("token-usage: {usage:?}");
handler.usage(usage);
}
}
_ => {}
}
}
@@ -323,7 +335,11 @@ struct EmbeddingsResBody {
embeddings: Vec<Vec<f32>>,
}
fn build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Result<Value> {
fn build_chat_completions_body(
data: ChatCompletionsData,
model: &Model,
prompt_cache: bool,
) -> Result<Value> {
let ChatCompletionsData {
mut messages,
temperature,
@@ -390,45 +406,112 @@ fn build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Resu
})]
}
MessageContent::ToolCalls(MessageContentToolCalls {
tool_results, text, ..
tool_results,
text,
sequence,
}) => {
let mut assistant_parts = vec![];
let mut user_parts = vec![];
if !text.is_empty() {
assistant_parts.push(json!({
"text": text,
}))
}
for tool_result in tool_results {
if let Some(round_text) = &tool_result.text {
// Unlike claude.rs, tool_result.thinking blocks are dropped:
// Converse reasoningContent replay is not implemented.
if !sequence {
let mut assistant_parts = vec![];
let mut user_parts = vec![];
for (index, tool_result) in tool_results.iter().enumerate() {
let round_text = if index == 0 && !text.is_empty() {
Some(text.as_str())
} else {
tool_result.text.as_deref()
};
if let Some(round_text) = round_text {
let round_text = strip_think_tag(round_text);
let round_text = round_text.trim();
if !round_text.is_empty() {
assistant_parts.push(json!({
"text": round_text,
}))
}
}
assistant_parts.push(json!({
"text": round_text,
}))
"toolUse": {
"toolUseId": tool_result.call.id,
"name": tool_result.call.name,
"input": tool_result.call.arguments,
}
}));
user_parts.push(json!({
"toolResult": {
"toolUseId": tool_result.call.id,
"content": [bedrock_tool_result_content(&tool_result.output)]
}
}));
}
assistant_parts.push(json!({
"toolUse": {
"toolUseId": tool_result.call.id,
"name": tool_result.call.name,
"input": tool_result.call.arguments,
vec![
json!({ "role": "assistant", "content": assistant_parts }),
json!({ "role": "user", "content": user_parts }),
]
} else {
// One pair per round: Claude can reuse tool_use IDs across API calls.
// A round boundary is detected by the presence of round text, but
// rounds where the model emitted only tool calls (no narration)
// carry no text marker. As a backstop, also split whenever a
// toolUseId would repeat within the current assistant message.
// Converse rejects duplicate toolUseIds in a single message.
let mut messages = vec![];
let mut assistant_parts: Vec<Value> = vec![];
let mut user_parts: Vec<Value> = vec![];
let mut chunk_ids: HashSet<&str> = HashSet::new();
for (index, tool_result) in tool_results.iter().enumerate() {
let id_collision = tool_result
.call
.id
.as_deref()
.is_some_and(|id| chunk_ids.contains(id));
if index > 0 && (tool_result.text.is_some() || id_collision) {
messages.push(
json!({ "role": "assistant", "content": assistant_parts }),
);
messages.push(json!({ "role": "user", "content": user_parts }));
assistant_parts = vec![];
user_parts = vec![];
chunk_ids.clear();
}
}));
user_parts.push(json!({
"toolResult": {
"toolUseId": tool_result.call.id,
"content": [bedrock_tool_result_content(&tool_result.output)]
if let Some(id) = tool_result.call.id.as_deref() {
chunk_ids.insert(id);
}
}));
let round_text = if index == 0 && !text.is_empty() {
Some(text.as_str())
} else {
tool_result.text.as_deref()
};
if let Some(round_text) = round_text {
let round_text = strip_think_tag(round_text);
let round_text = round_text.trim();
if !round_text.is_empty() {
assistant_parts.push(json!({
"text": round_text,
}))
}
}
assistant_parts.push(json!({
"toolUse": {
"toolUseId": tool_result.call.id,
"name": tool_result.call.name,
"input": tool_result.call.arguments,
}
}));
user_parts.push(json!({
"toolResult": {
"toolUseId": tool_result.call.id,
"content": [bedrock_tool_result_content(&tool_result.output)]
}
}));
}
if !assistant_parts.is_empty() {
messages
.push(json!({ "role": "assistant", "content": assistant_parts }));
messages.push(json!({ "role": "user", "content": user_parts }));
}
messages
}
vec![
json!({
"role": "assistant",
"content": assistant_parts,
}),
json!({
"role": "user",
"content": user_parts,
}),
]
}
}
})
@@ -484,9 +567,70 @@ fn build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Resu
"tools": tools,
})
}
if prompt_cache && is_anthropic_model(model) {
apply_prompt_cache_points(&mut body);
}
Ok(body)
}
/// Converse rejects cachePoint blocks for models that don't support them
/// (e.g. openai.gpt-*), so caching is only applied to Anthropic models.
fn is_anthropic_model(model: &Model) -> bool {
let name = model.real_name();
name.contains("anthropic.") || name.contains("claude")
}
/// Converse counterpart of the claude client's cache_control breakpoints:
/// a cachePoint after the tools, the system prompt, and the last two user
/// messages. See `apply_prompt_cache_breakpoints` in claude.rs for why the
/// last two.
fn apply_prompt_cache_points(body: &mut Value) {
let cache_point = || json!({ "cachePoint": { "type": "default" } });
if let Some(system) = body.get_mut("system").and_then(|v| v.as_array_mut()) {
system.push(cache_point());
}
if let Some(tools) = body
.get_mut("toolConfig")
.and_then(|v| v.get_mut("tools"))
.and_then(|v| v.as_array_mut())
{
tools.push(cache_point());
}
let Some(messages) = body.get_mut("messages").and_then(|v| v.as_array_mut()) else {
return;
};
let mut remaining = 2;
for message in messages.iter_mut().rev() {
if remaining == 0 {
break;
}
if message["role"] != "user" {
continue;
}
// Converse rejects a cachePoint with no preceding block, so an empty
// content array (a tool-call pair with no results) passes the
// breakpoint to an earlier user message.
if let Some(content) = message["content"].as_array_mut()
&& !content.is_empty()
{
content.push(cache_point());
remaining -= 1;
}
}
}
fn bedrock_parse_usage(usage: &Value) -> Option<TokenUsage> {
if !usage.is_object() {
return None;
}
Some(TokenUsage {
input_tokens: usage["inputTokens"].as_u64(),
output_tokens: usage["outputTokens"].as_u64(),
cache_creation_input_tokens: usage["cacheWriteInputTokens"].as_u64(),
cache_read_input_tokens: usage["cacheReadInputTokens"].as_u64(),
})
}
fn bedrock_tool_result_content(output: &Value) -> Value {
match output {
Value::Object(_) => json!({ "json": output }),
@@ -539,6 +683,7 @@ fn extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
let output = ChatCompletionsOutput {
text,
tool_calls,
usage: bedrock_parse_usage(&data["usage"]),
..Default::default()
};
Ok(output)
@@ -661,7 +806,7 @@ fn gen_signing_key(key: &str, date_stamp: &str, region: &str, service: &str) ->
#[cfg(test)]
mod tests {
use super::*;
use crate::function::{ToolCall, ToolResult};
use crate::function::{FunctionDeclaration, ToolCall, ToolResult};
fn tool_result(output: Value) -> ToolResult {
ToolResult {
@@ -676,14 +821,23 @@ mod tests {
}
}
fn tool_result_content_block(output: Value) -> Value {
fn sequence_tool_result(id: &str, text: Option<&str>) -> ToolResult {
ToolResult {
call: ToolCall::new("fs_read".into(), json!({"path": "x"}), Some(id.into())),
output: json!("ok"),
text: text.map(|t| t.to_string()),
thinking: vec![],
}
}
fn build_body(tool_results: Vec<ToolResult>) -> Value {
let data = ChatCompletionsData {
messages: vec![
Message::new(MessageRole::User, MessageContent::Text("hello".to_string())),
Message::new(
MessageRole::Assistant,
MessageContent::ToolCalls(MessageContentToolCalls {
tool_results: vec![tool_result(output)],
tool_results,
text: String::new(),
sequence: true,
}),
@@ -695,10 +849,178 @@ mod tests {
functions: None,
stream: false,
};
let body = build_chat_completions_body(data, &Model::new("bedrock", "test")).unwrap();
build_chat_completions_body(data, &Model::new("bedrock", "test"), false).unwrap()
}
fn tool_result_content_block(output: Value) -> Value {
let body = build_body(vec![tool_result(output)]);
body["messages"][2]["content"][0]["toolResult"]["content"][0].clone()
}
fn assert_unique_tool_use_ids_per_message(body: &Value) {
for message in body["messages"].as_array().unwrap() {
let Some(content) = message["content"].as_array() else {
continue;
};
let mut seen = HashSet::new();
for block in content {
if let Some(id) = block["toolUse"]["toolUseId"].as_str() {
assert!(
seen.insert(id.to_string()),
"duplicate toolUseId `{id}` within a single assistant message: {message}"
);
}
}
}
}
fn cache_fixture() -> ChatCompletionsData {
ChatCompletionsData {
messages: vec![
Message::new(MessageRole::System, MessageContent::Text("sys".to_string())),
Message::new(MessageRole::User, MessageContent::Text("hello".to_string())),
Message::new(
MessageRole::Assistant,
MessageContent::Text("hi".to_string()),
),
Message::new(MessageRole::User, MessageContent::Text("again".to_string())),
],
temperature: None,
top_p: None,
reasoning_effort: None,
functions: Some(vec![FunctionDeclaration {
name: "fs_read".to_string(),
description: "Read a file".to_string(),
parameters: Default::default(),
agent: false,
}]),
stream: false,
}
}
#[test]
fn test_bedrock_prompt_cache_inserts_cache_points_for_anthropic_models() {
let body = build_chat_completions_body(
cache_fixture(),
&Model::new("bedrock", "anthropic.claude-sonnet-4-5-20250929-v1:0"),
true,
)
.unwrap();
let cache_point = json!({ "cachePoint": { "type": "default" } });
assert_eq!(
body["system"].as_array().unwrap().last(),
Some(&cache_point),
"body: {body}"
);
assert_eq!(
body["toolConfig"]["tools"].as_array().unwrap().last(),
Some(&cache_point),
"body: {body}"
);
let messages = body["messages"].as_array().unwrap();
for idx in [0, 2] {
assert_eq!(
messages[idx]["content"].as_array().unwrap().last(),
Some(&cache_point),
"body: {body}"
);
}
assert!(
messages[1]["content"]
.as_array()
.unwrap()
.iter()
.all(|block| block.get("cachePoint").is_none()),
"assistant messages must not get cache points: {body}"
);
}
#[test]
fn test_bedrock_prompt_cache_skips_non_anthropic_models() {
let body = build_chat_completions_body(
cache_fixture(),
&Model::new("bedrock", "openai.gpt-oss-120b-1:0"),
true,
)
.unwrap();
assert!(!body.to_string().contains("cachePoint"), "body: {body}");
}
/// Converse rejects a cachePoint with no preceding block, so a user
/// message with an empty content array must be skipped and the
/// breakpoint placed on an earlier user message instead.
#[test]
fn test_bedrock_prompt_cache_skips_empty_user_content() {
let mut body = json!({
"messages": [
{ "role": "user", "content": [ { "text": "hello" } ] },
{ "role": "assistant", "content": [ { "text": "hi" } ] },
{ "role": "user", "content": [ { "text": "again" } ] },
{ "role": "assistant", "content": [ { "text": "calling" } ] },
{ "role": "user", "content": [] },
],
});
apply_prompt_cache_points(&mut body);
let cache_point = json!({ "cachePoint": { "type": "default" } });
let messages = body["messages"].as_array().unwrap();
assert!(
messages[4]["content"].as_array().unwrap().is_empty(),
"empty user content must not get a cachePoint: {body}"
);
for idx in [0, 2] {
assert_eq!(
messages[idx]["content"].as_array().unwrap().last(),
Some(&cache_point),
"body: {body}"
);
}
}
#[test]
fn sequence_splits_on_round_text() {
let body = build_body(vec![
sequence_tool_result("toolu_A", None),
sequence_tool_result("toolu_B", None),
sequence_tool_result("toolu_C", Some("running another tool")),
]);
let messages = body["messages"].as_array().unwrap();
assert_eq!(messages.len(), 5, "body: {body}");
assert_unique_tool_use_ids_per_message(&body);
}
#[test]
fn sequence_splits_on_reused_id_in_textless_round() {
let body = build_body(vec![
sequence_tool_result("toolu_A", None),
sequence_tool_result("toolu_B", None),
sequence_tool_result("toolu_A", None),
]);
let messages = body["messages"].as_array().unwrap();
assert_eq!(messages.len(), 5, "body: {body}");
assert_unique_tool_use_ids_per_message(&body);
}
#[test]
fn sequence_keeps_textless_rounds_merged_when_ids_are_unique() {
let body = build_body(vec![
sequence_tool_result("toolu_A", None),
sequence_tool_result("toolu_B", None),
sequence_tool_result("toolu_C", None),
]);
let messages = body["messages"].as_array().unwrap();
assert_eq!(messages.len(), 3, "body: {body}");
assert_unique_tool_use_ids_per_message(&body);
}
#[test]
fn test_bedrock_tool_result_object_uses_json_block() {
assert_eq!(
src/client/claude.rs
+230 -19
View File
@@ -24,6 +24,7 @@ pub struct ClaudeConfig {
pub auth: Option<String>,
#[serde(default)]
pub models: Vec<ModelData>,
pub prompt_cache: Option<bool>,
pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>,
}
@@ -75,7 +76,11 @@ async fn prepare_chat_completions(
.unwrap_or_else(|_| API_BASE.to_string());
let url = format!("{}/messages", api_base.trim_end_matches('/'));
let body = claude_build_chat_completions_body(data, &self_.model)?;
let body = claude_build_chat_completions_body(
data,
&self_.model,
self_.config.prompt_cache.unwrap_or(true),
)?;
let mut request_data = RequestData::new(url, body);
@@ -250,6 +255,18 @@ pub async fn claude_chat_completions_streaming(
function_arguments.clear();
}
}
"message_start" => {
if let Some(usage) = claude_parse_usage(&data["message"]["usage"]) {
debug!("token-usage: {usage:?}");
handler.usage(usage);
}
}
"message_delta" => {
if let Some(usage) = claude_parse_usage(&data["usage"]) {
debug!("token-usage: {usage:?}");
handler.usage(usage);
}
}
_ => {}
}
}
@@ -262,6 +279,7 @@ pub async fn claude_chat_completions_streaming(
pub fn claude_build_chat_completions_body(
data: ChatCompletionsData,
model: &Model,
prompt_cache: bool,
) -> Result<Value> {
let ChatCompletionsData {
mut messages,
@@ -487,9 +505,65 @@ pub fn claude_build_chat_completions_body(
})
.collect();
}
if prompt_cache {
apply_prompt_cache_breakpoints(&mut body);
}
Ok(body)
}
/// Anthropic allows at most 4 cache_control breakpoints; all four are spent
/// here: the last tool, the system prompt, and the last two user messages.
/// Two moving breakpoints guarantee a cache read every turn: turn N+1's
/// second-to-last breakpoint sits exactly where turn N's last one was.
fn apply_prompt_cache_breakpoints(body: &mut Value) {
if let Some(text) = body["system"].as_str().map(str::to_string) {
body["system"] = json!([{
"type": "text",
"text": text,
"cache_control": { "type": "ephemeral" },
}]);
}
if let Some(tool) = body
.get_mut("tools")
.and_then(|v| v.as_array_mut())
.and_then(|v| v.last_mut())
{
tool["cache_control"] = json!({ "type": "ephemeral" });
}
let Some(messages) = body.get_mut("messages").and_then(|v| v.as_array_mut()) else {
return;
};
let mut remaining = 2;
for message in messages.iter_mut().rev() {
if remaining == 0 {
break;
}
if message["role"] != "user" {
continue;
}
if let Some(text) = message["content"].as_str().map(str::to_string) {
message["content"] = json!([{ "type": "text", "text": text }]);
}
if let Some(block) = message["content"].as_array_mut().and_then(|v| v.last_mut()) {
block["cache_control"] = json!({ "type": "ephemeral" });
remaining -= 1;
}
}
}
fn claude_parse_usage(usage: &Value) -> Option<TokenUsage> {
if !usage.is_object() {
return None;
}
Some(TokenUsage {
input_tokens: usage["input_tokens"].as_u64(),
output_tokens: usage["output_tokens"].as_u64(),
cache_creation_input_tokens: usage["cache_creation_input_tokens"].as_u64(),
cache_read_input_tokens: usage["cache_read_input_tokens"].as_u64(),
})
}
pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOutput> {
let mut text = String::new();
let mut tool_calls = vec![];
@@ -545,6 +619,7 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
text: text.to_string(),
tool_calls,
thinking,
usage: claude_parse_usage(&data["usage"]),
};
Ok(output)
}
@@ -552,7 +627,7 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result<ChatCompletionsOu
#[cfg(test)]
mod tests {
use super::*;
use crate::function::{ToolCall, ToolResult};
use crate::function::{FunctionDeclaration, ToolCall, ToolResult};
fn tool_result(id: &str, text: Option<&str>) -> ToolResult {
ToolResult {
@@ -563,7 +638,7 @@ mod tests {
}
}
fn build_body(tool_results: Vec<ToolResult>) -> Value {
fn build_body(tool_results: Vec<ToolResult>, prompt_cache: bool) -> Value {
let data = ChatCompletionsData {
messages: vec![
Message::new(MessageRole::User, MessageContent::Text("hello".to_string())),
@@ -582,7 +657,44 @@ mod tests {
functions: None,
stream: false,
};
claude_build_chat_completions_body(data, &Model::new("claude", "claude-test")).unwrap()
claude_build_chat_completions_body(data, &Model::new("claude", "claude-test"), prompt_cache)
.unwrap()
}
fn multi_turn_data(functions: Option<Vec<FunctionDeclaration>>) -> ChatCompletionsData {
ChatCompletionsData {
messages: vec![
Message::new(MessageRole::System, MessageContent::Text("sys".to_string())),
Message::new(MessageRole::User, MessageContent::Text("first".to_string())),
Message::new(
MessageRole::Assistant,
MessageContent::Text("reply".to_string()),
),
Message::new(
MessageRole::User,
MessageContent::Text("second".to_string()),
),
Message::new(
MessageRole::Assistant,
MessageContent::Text("reply2".to_string()),
),
Message::new(MessageRole::User, MessageContent::Text("third".to_string())),
],
temperature: None,
top_p: None,
reasoning_effort: None,
functions,
stream: false,
}
}
fn function_declaration(name: &str) -> FunctionDeclaration {
FunctionDeclaration {
name: name.to_string(),
description: format!("{name} description"),
parameters: Default::default(),
agent: false,
}
}
fn assert_unique_tool_use_ids_per_message(body: &Value) {
@@ -605,11 +717,14 @@ mod tests {
#[test]
fn sequence_splits_on_round_text() {
let body = build_body(vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_C", Some("running another tool")),
]);
let body = build_body(
vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_C", Some("running another tool")),
],
false,
);
let messages = body["messages"].as_array().unwrap();
@@ -619,11 +734,14 @@ mod tests {
#[test]
fn sequence_splits_on_reused_id_in_textless_round() {
let body = build_body(vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_A", None),
]);
let body = build_body(
vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_A", None),
],
false,
);
let messages = body["messages"].as_array().unwrap();
@@ -633,15 +751,108 @@ mod tests {
#[test]
fn sequence_keeps_textless_rounds_merged_when_ids_are_unique() {
let body = build_body(vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_C", None),
]);
let body = build_body(
vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_C", None),
],
false,
);
let messages = body["messages"].as_array().unwrap();
assert_eq!(messages.len(), 3, "body: {body}");
assert_unique_tool_use_ids_per_message(&body);
}
#[test]
fn prompt_cache_places_breakpoints() {
let functions = vec![function_declaration("a"), function_declaration("b")];
let body = claude_build_chat_completions_body(
multi_turn_data(Some(functions)),
&Model::new("claude", "claude-test"),
true,
)
.unwrap();
let cache_control = json!({ "type": "ephemeral" });
let tools = body["tools"].as_array().unwrap();
assert!(tools[0].get("cache_control").is_none(), "body: {body}");
assert_eq!(tools[1]["cache_control"], cache_control);
let system = body["system"].as_array().unwrap();
assert_eq!(system.last().unwrap()["text"], "sys");
assert_eq!(system.last().unwrap()["cache_control"], cache_control);
let messages = body["messages"].as_array().unwrap();
assert!(
messages[0]["content"].is_string(),
"older user messages stay untouched: {body}"
);
for idx in [2, 4] {
let blocks = messages[idx]["content"].as_array().unwrap();
let last = blocks.last().unwrap();
assert_eq!(last["type"], "text", "body: {body}");
assert_eq!(last["cache_control"], cache_control, "body: {body}");
}
}
#[test]
fn prompt_cache_marks_last_tool_result_block() {
let body = build_body(
vec![tool_result("toolu_A", None), tool_result("toolu_B", None)],
true,
);
assert!(body.get("system").is_none(), "body: {body}");
assert!(body.get("tools").is_none(), "body: {body}");
let cache_control = json!({ "type": "ephemeral" });
let messages = body["messages"].as_array().unwrap();
let last_blocks = messages.last().unwrap()["content"].as_array().unwrap();
let last_block = last_blocks.last().unwrap();
assert_eq!(last_block["type"], "tool_result", "body: {body}");
assert_eq!(last_block["cache_control"], cache_control);
assert!(last_blocks[0].get("cache_control").is_none());
let first_blocks = messages[0]["content"].as_array().unwrap();
assert_eq!(first_blocks.last().unwrap()["type"], "text");
assert_eq!(first_blocks.last().unwrap()["cache_control"], cache_control);
}
#[test]
fn prompt_cache_disabled_leaves_body_unchanged() {
let functions = vec![function_declaration("a")];
let body = claude_build_chat_completions_body(
multi_turn_data(Some(functions)),
&Model::new("claude", "claude-test"),
false,
)
.unwrap();
assert!(body["system"].is_string(), "body: {body}");
assert!(!body.to_string().contains("cache_control"), "body: {body}");
}
#[test]
fn extract_populates_usage() {
let data = json!({
"content": [{ "type": "text", "text": "hi" }],
"usage": {
"input_tokens": 10,
"output_tokens": 5,
"cache_creation_input_tokens": 100,
"cache_read_input_tokens": 200,
}
});
let output = claude_extract_chat_completions(&data).unwrap();
let usage = output.usage.unwrap();
assert_eq!(usage.input_tokens, Some(10));
assert_eq!(usage.output_tokens, Some(5));
assert_eq!(usage.cache_creation_input_tokens, Some(100));
assert_eq!(usage.cache_read_input_tokens, Some(200));
}
}
src/client/common.rs
+100 -8
View File
@@ -17,7 +17,7 @@ use inquire::{
MultiSelect, Select, Text, list_option::ListOption, required, validator::Validation,
};
use reqwest::{Client as ReqwestClient, RequestBuilder};
use serde::Deserialize;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::sync::LazyLock;
use std::time::Duration;
@@ -357,6 +357,7 @@ pub struct ChatCompletionsOutput {
pub text: String,
pub tool_calls: Vec<ToolCall>,
pub thinking: Vec<ThinkingBlock>,
pub usage: Option<TokenUsage>,
}
impl ChatCompletionsOutput {
@@ -368,6 +369,41 @@ impl ChatCompletionsOutput {
}
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct TokenUsage {
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub cache_creation_input_tokens: Option<u64>,
pub cache_read_input_tokens: Option<u64>,
}
impl TokenUsage {
pub fn is_empty(&self) -> bool {
self.input_tokens.is_none()
&& self.output_tokens.is_none()
&& self.cache_creation_input_tokens.is_none()
&& self.cache_read_input_tokens.is_none()
}
pub fn accumulate(&mut self, other: &TokenUsage) {
fn add(acc: &mut Option<u64>, value: Option<u64>) {
if let Some(value) = value {
*acc = Some(acc.unwrap_or(0).saturating_add(value));
}
}
add(&mut self.input_tokens, other.input_tokens);
add(&mut self.output_tokens, other.output_tokens);
add(
&mut self.cache_creation_input_tokens,
other.cache_creation_input_tokens,
);
add(
&mut self.cache_read_input_tokens,
other.cache_read_input_tokens,
);
}
}
#[derive(Debug)]
pub struct EmbeddingsData {
pub texts: Vec<String>,
@@ -513,7 +549,7 @@ pub async fn call_chat_completions(
mut text,
tool_calls,
thinking,
..
usage,
} = ret;
if !text.is_empty() {
if extract_code {
@@ -523,6 +559,7 @@ pub async fn call_chat_completions(
ctx.app.config.print_markdown(&text)?;
}
}
ctx.record_token_usage(usage);
finish_completion(ctx, text, tool_calls, thinking).await
}
Err(err) => Err(err),
@@ -550,14 +587,18 @@ pub async fn call_chat_completions_streaming(
let aborted_ctrlc = handler.abort().aborted_ctrlc();
let aborted_ctrld = handler.abort().aborted_ctrld();
// Record usage before any early return: the provider bills an interrupted
// call's input tokens (reported in message_start), so abort and error
// paths must still count toward the session totals.
let (text, tool_calls, thinking, usage) = handler.take();
ctx.record_token_usage(usage);
if aborted_ctrld {
bail!("Aborted.");
}
render_ret?;
let (text, tool_calls, thinking) = handler.take();
if aborted_ctrlc {
if !ctx.working_mode.is_repl() || ctx.session.is_none() {
bail!("Aborted.");
@@ -618,14 +659,16 @@ pub async fn call_chat_completions_streaming_quiet(
ret = client.chat_completions(input.clone()) => ret?,
_ = wait_abort_signal(&abort_signal) => bail!("Aborted."),
};
if abort_signal.aborted() {
bail!("Aborted.");
}
let ChatCompletionsOutput {
text,
tool_calls,
thinking,
usage,
} = output;
ctx.record_token_usage(usage);
if abort_signal.aborted() {
bail!("Aborted.");
}
return finish_completion(ctx, text, tool_calls, thinking).await;
}
@@ -637,15 +680,19 @@ pub async fn call_chat_completions_streaming_quiet(
let send_ret = client.chat_completions_streaming(input, &mut handler).await;
let (text, tool_calls, thinking, usage) = handler.take();
ctx.record_token_usage(usage);
if abort_signal.aborted() {
bail!("Aborted.");
}
send_ret?;
let (text, tool_calls, thinking) = handler.take();
finish_completion(ctx, text, tool_calls, thinking).await
}
/// Callers record token usage themselves before calling, so that abort and
/// error paths that never reach here still count exactly once per API call.
async fn finish_completion(
ctx: &mut RequestContext,
text: String,
@@ -909,6 +956,50 @@ mod tests {
catch_error(data, status).unwrap_err().to_string()
}
#[test]
fn test_token_usage_is_empty_treats_some_zero_as_non_empty() {
assert!(TokenUsage::default().is_empty());
let usage = TokenUsage {
input_tokens: Some(0),
..Default::default()
};
assert!(!usage.is_empty());
}
#[test]
fn test_token_usage_accumulate_merges_per_field() {
let mut total = TokenUsage::default();
total.accumulate(&TokenUsage {
input_tokens: Some(100),
output_tokens: None,
cache_creation_input_tokens: Some(50),
cache_read_input_tokens: None,
});
total.accumulate(&TokenUsage {
input_tokens: Some(20),
output_tokens: Some(7),
cache_creation_input_tokens: None,
cache_read_input_tokens: Some(90),
});
assert_eq!(total.input_tokens, Some(120));
assert_eq!(total.output_tokens, Some(7));
assert_eq!(total.cache_creation_input_tokens, Some(50));
assert_eq!(total.cache_read_input_tokens, Some(90));
}
#[test]
fn test_token_usage_accumulate_saturates() {
let mut total = TokenUsage {
input_tokens: Some(u64::MAX),
..Default::default()
};
total.accumulate(&TokenUsage {
input_tokens: Some(1),
..Default::default()
});
assert_eq!(total.input_tokens, Some(u64::MAX));
}
#[test]
fn test_catch_error_display_json_with_type() {
let data = json!({"error": {"type": "invalid_request_error", "message": "Bad request"}});
@@ -1141,6 +1232,7 @@ mod tests {
text: paced_text(),
tool_calls: vec![paced_tool_call()],
thinking: vec![paced_thinking()],
usage: None,
})
}
src/client/stream.rs
+61 -5
View File
@@ -1,4 +1,4 @@
use super::{ApiStatusError, ThinkingBlock, ToolCall, catch_error};
use super::{ApiStatusError, ThinkingBlock, TokenUsage, ToolCall, catch_error};
use crate::utils::AbortSignal;
use anyhow::{Context, Result, anyhow, bail};
@@ -14,6 +14,7 @@ pub struct SseHandler {
buffer: String,
tool_calls: Vec<ToolCall>,
thinking: Vec<ThinkingBlock>,
usage: Option<TokenUsage>,
last_tool_calls: Vec<ToolCall>,
max_call_repeats: usize,
call_repeat_chain_len: usize,
@@ -29,6 +30,7 @@ impl SseHandler {
buffer: String::new(),
tool_calls: Vec::new(),
thinking: Vec::new(),
usage: None,
last_tool_calls: Vec::new(),
max_call_repeats: 2,
call_repeat_chain_len: 3,
@@ -185,6 +187,26 @@ impl SseHandler {
self.thinking.push(block);
}
/// Merges a usage update into the accumulator; later non-None fields win.
/// Claude's `message_start` carries input/cache counts while
/// `message_delta` carries the output count, so a single event never has
/// the full picture.
pub fn usage(&mut self, update: TokenUsage) {
let usage = self.usage.get_or_insert_with(TokenUsage::default);
if update.input_tokens.is_some() {
usage.input_tokens = update.input_tokens;
}
if update.output_tokens.is_some() {
usage.output_tokens = update.output_tokens;
}
if update.cache_creation_input_tokens.is_some() {
usage.cache_creation_input_tokens = update.cache_creation_input_tokens;
}
if update.cache_read_input_tokens.is_some() {
usage.cache_read_input_tokens = update.cache_read_input_tokens;
}
}
/// Whether any output (text, tool calls, or thinking blocks) has been
/// accumulated. `Client::chat_completions_streaming` gates its 401 retry
/// on this: content already streamed to the user would be rendered a
@@ -202,14 +224,22 @@ impl SseHandler {
&self.last_tool_calls
}
pub fn take(self) -> (String, Vec<ToolCall>, Vec<ThinkingBlock>) {
pub fn take(
self,
) -> (
String,
Vec<ToolCall>,
Vec<ThinkingBlock>,
Option<TokenUsage>,
) {
let Self {
buffer,
tool_calls,
thinking,
usage,
..
} = self;
(buffer, tool_calls, thinking)
(buffer, tool_calls, thinking, usage)
}
}
@@ -455,7 +485,7 @@ mod tests {
handler.tool_call(call.clone()).unwrap();
}
let (_, calls, _) = handler.take();
let (_, calls, _, _) = handler.take();
assert_eq!(calls.len(), 5);
}
@@ -514,6 +544,32 @@ mod tests {
assert!(handler.has_received_content());
}
/// Claude splits usage across events: `message_start` carries the input
/// and cache counts with a provisional output count, `message_delta`
/// carries only the final output count. Later non-None fields must win.
#[test]
fn test_usage_merges_later_non_none_fields() {
let (mut handler, _rx) = new_handler();
handler.usage(TokenUsage {
input_tokens: Some(100),
output_tokens: Some(1),
cache_creation_input_tokens: Some(200),
cache_read_input_tokens: Some(300),
});
handler.usage(TokenUsage {
output_tokens: Some(42),
..Default::default()
});
let (_, _, _, usage) = handler.take();
let usage = usage.unwrap();
assert_eq!(usage.input_tokens, Some(100));
assert_eq!(usage.output_tokens, Some(42));
assert_eq!(usage.cache_creation_input_tokens, Some(200));
assert_eq!(usage.cache_read_input_tokens, Some(300));
}
/// A silent handler is the quiet transport's whole output gate: text,
/// tool calls, and thinking still accumulate for `take()`, but the
/// receiver sees no `Text` events, only the trailing `Done`.
@@ -540,7 +596,7 @@ mod tests {
assert!(matches!(rx.try_recv(), Ok(SseEvent::Done)));
assert!(rx.try_recv().is_err(), "no events beyond Done");
let (text, calls, thinking) = handler.take();
let (text, calls, thinking, _) = handler.take();
assert_eq!(text, "hello");
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "lookup");
src/client/vertexai.rs
+6 -1
View File
@@ -18,6 +18,7 @@ pub struct VertexAIConfig {
pub adc_file: Option<String>,
#[serde(default)]
pub models: Vec<ModelData>,
pub prompt_cache: Option<bool>,
pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>,
}
@@ -137,7 +138,11 @@ fn prepare_chat_completions(
let body = match model_category {
ModelCategory::Gemini => gemini_build_chat_completions_body(data, &self_.model)?,
ModelCategory::Claude => {
let mut body = claude_build_chat_completions_body(data, &self_.model)?;
let mut body = claude_build_chat_completions_body(
data,
&self_.model,
self_.config.prompt_cache.unwrap_or(true),
)?;
if let Some(body_obj) = body.as_object_mut() {
body_obj.remove("model");
}