diff --git a/assets/config-template.yaml b/assets/config-template.yaml index 0ab0ed8..7e566f3 100644 --- a/assets/config-template.yaml +++ b/assets/config-template.yaml @@ -240,7 +240,10 @@ theme: null # null = the built-in dark theme; set to `light left_prompt: null # If null, uses the built-in default: # '{color.red}{model}){color.green}{?session {?agent {agent}>}{session}{?role /}}{!session {?agent {agent}>}}{role}{?rag @{rag}}{color.cyan}{?session )}{!session >}{color.reset} ' right_prompt: null # If null, uses the built-in default: - # '{color.cyan}{?reasoning_effort [{reasoning_effort}] }{color.purple}{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}' + # '{color.cyan}{?reasoning_effort [{reasoning_effort}] }{color.purple}{?last_cache_read_tokens ⚡{last_cache_read_tokens} }{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}' + # Token-usage variables (set only when the provider reports usage, so guard them with {?var ...}): + # last_input_tokens, last_output_tokens, last_cache_read_tokens, last_cache_creation_tokens (most recent API call) + # total_input_tokens, total_output_tokens, total_cache_read_tokens, total_cache_creation_tokens (session cumulative) # ---- Miscellaneous ---- user_agent: null # Set User-Agent HTTP header, use `auto` for coyote/ @@ -282,4 +285,7 @@ sync_models_ttl_hours: 24 # Hours before the synced model catalog is cons # proxy: socks5://127.0.0.1:1080 # Set proxy # connect_timeout: 10 # Set timeout in seconds for connect to api # read_timeout: 300 # Set timeout in seconds for a read stall (no bytes received); 0 disables (default: 300) +# +# The claude, bedrock, and vertexai clients also support: +# prompt_cache: true # Enable Anthropic prompt caching (default: true) __CLIENTS_BLOCK__ diff --git a/config.example.yaml b/config.example.yaml index f4bd891..c48dd58 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -36,10 +36,13 @@ sync_models_ttl_hours: 24 # Hours before the synced model catalog is cons # ---- REPL Prompt ---- # Custom REPL left/right prompts; see the [REPL Prompt Documentation](https://github.com/Dark-Alex-17/coyote/wiki/REPL-Prompt) for more information +# Token-usage variables (set only when the provider reports usage, so guard them with {?var ...}): +# last_input_tokens, last_output_tokens, last_cache_read_tokens, last_cache_creation_tokens (most recent API call) +# total_input_tokens, total_output_tokens, total_cache_read_tokens, total_cache_creation_tokens (session cumulative) left_prompt: '{color.red}{model}){color.green}{?session {?agent {agent}>}{session}{?role /}}{!session {?agent {agent}>}}{role}{?rag @{rag}}{color.cyan}{?session )}{!session >}{color.reset} ' right_prompt: - '{color.cyan}{?reasoning_effort [{reasoning_effort}] }{color.purple}{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}' + '{color.cyan}{?reasoning_effort [{reasoning_effort}] }{color.purple}{?last_cache_read_tokens ⚡{last_cache_read_tokens} }{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}' # ---- Vault ---- # See the [Vault documentation](https://github.com/Dark-Alex-17/coyote/wiki/Vault) for more information on the Coyote vault. @@ -380,6 +383,7 @@ clients: api_key: '{{ANTHROPIC_API_KEY}}' # You can either hard-code or inject secrets from the Coyote vault auth: null # When set to 'oauth', Coyote will use OAuth instead of an API key # Authenticate with `coyote --authenticate` or `.authenticate` in the REPL + prompt_cache: true # Optional - enable Anthropic prompt caching (default: true) # See https://docs.mistral.ai/ - type: openai-compatible @@ -461,6 +465,7 @@ clients: # Run `gcloud auth application-default login` to initialize the ADC file # see https://cloud.google.com/docs/authentication/external/set-up-adc for more information adc_file: /application_default_credentials.json # Optional + prompt_cache: true # Optional - enable Anthropic prompt caching for Claude models (default: true) patch: chat_completions: 'gemini-.*': @@ -481,6 +486,7 @@ clients: secret_access_key: '{{AWS_SECRET_ACCESS_KEY}}' # You can either hard-code or inject secrets from the Coyote vault region: xxx session_token: xxx # Optional, only needed for temporary credentials + prompt_cache: true # Optional - enable Anthropic prompt caching for Anthropic models (default: true) # See https://developers.cloudflare.com/workers-ai/ - type: openai-compatible diff --git a/src/client/bedrock.rs b/src/client/bedrock.rs index 5de0328..73781da 100644 --- a/src/client/bedrock.rs +++ b/src/client/bedrock.rs @@ -1,3 +1,4 @@ +use std::collections::HashSet; use std::mem; use super::*; @@ -24,6 +25,7 @@ pub struct BedrockConfig { pub session_token: Option, #[serde(default)] pub models: Vec, + pub prompt_cache: Option, pub patch: Option, pub extra: Option, } @@ -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>, } -fn build_chat_completions_body(data: ChatCompletionsData, model: &Model) -> Result { +fn build_chat_completions_body( + data: ChatCompletionsData, + model: &Model, + prompt_cache: bool, +) -> Result { 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 = vec![]; + let mut user_parts: Vec = 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 { + 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 { 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) -> 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!( diff --git a/src/client/claude.rs b/src/client/claude.rs index 7e26772..019b379 100644 --- a/src/client/claude.rs +++ b/src/client/claude.rs @@ -24,6 +24,7 @@ pub struct ClaudeConfig { pub auth: Option, #[serde(default)] pub models: Vec, + pub prompt_cache: Option, pub patch: Option, pub extra: Option, } @@ -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 { 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 { + 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 { let mut text = String::new(); let mut tool_calls = vec![]; @@ -545,6 +619,7 @@ pub fn claude_extract_chat_completions(data: &Value) -> Result Result) -> ToolResult { ToolResult { @@ -563,7 +638,7 @@ mod tests { } } - fn build_body(tool_results: Vec) -> Value { + fn build_body(tool_results: Vec, 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>) -> 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)); + } } diff --git a/src/client/common.rs b/src/client/common.rs index ac51ee6..e0c1965 100644 --- a/src/client/common.rs +++ b/src/client/common.rs @@ -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, pub thinking: Vec, + pub usage: Option, } impl ChatCompletionsOutput { @@ -368,6 +369,41 @@ impl ChatCompletionsOutput { } } +#[derive(Debug, Clone, Default, Deserialize, Serialize)] +pub struct TokenUsage { + pub input_tokens: Option, + pub output_tokens: Option, + pub cache_creation_input_tokens: Option, + pub cache_read_input_tokens: Option, +} + +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, value: Option) { + 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, @@ -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, }) } diff --git a/src/client/stream.rs b/src/client/stream.rs index 8d37aa4..2d9d3ab 100644 --- a/src/client/stream.rs +++ b/src/client/stream.rs @@ -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, thinking: Vec, + usage: Option, last_tool_calls: Vec, 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, Vec) { + pub fn take( + self, + ) -> ( + String, + Vec, + Vec, + Option, + ) { 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"); diff --git a/src/client/vertexai.rs b/src/client/vertexai.rs index e3705fb..3f1f6f5 100644 --- a/src/client/vertexai.rs +++ b/src/client/vertexai.rs @@ -18,6 +18,7 @@ pub struct VertexAIConfig { pub adc_file: Option, #[serde(default)] pub models: Vec, + pub prompt_cache: Option, pub patch: Option, pub extra: Option, } @@ -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"); } diff --git a/src/config/mod.rs b/src/config/mod.rs index 1016014..b13aa06 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -217,7 +217,8 @@ Rules: const SUMMARY_CONTEXT_PROMPT: &str = "This is a summary of the chat history as a recap: "; const LEFT_PROMPT: &str = "{color.red}{model}){color.green}{?session {?agent {agent}>}{session}{?role /}}{!session {?agent {agent}>}}{role}{?rag @{rag}}{color.cyan}{?session )}{!session >}{color.reset} "; -const RIGHT_PROMPT: &str = "{color.cyan}{?reasoning_effort [{reasoning_effort}] }{color.purple}{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}"; +// `{?var}` treats "0" as falsy, so a reported-but-zero cache read intentionally hides the ⚡ segment. +const RIGHT_PROMPT: &str = "{color.cyan}{?reasoning_effort [{reasoning_effort}] }{color.purple}{?last_cache_read_tokens ⚡{last_cache_read_tokens} }{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}"; static EDITOR: OnceLock> = OnceLock::new(); diff --git a/src/config/request_context.rs b/src/config/request_context.rs index b669ed6..c208602 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -16,7 +16,7 @@ use super::{ list_agents_with_descriptions, memory, paths, }; use super::{MessageContentToolCalls, prompts}; -use crate::client::{Model, ModelType, list_models}; +use crate::client::{Model, ModelType, TokenUsage, list_models}; use crate::function::{ FunctionDeclaration, Functions, ToolCallTracker, ToolResult, agents::AGENT_FUNCTION_PREFIX, @@ -308,6 +308,10 @@ pub struct RequestContext { pub last_message: Option, + /// Usage reported by the most recent API call; reset on every completion + /// so providers that report nothing never show stale numbers. + pub last_token_usage: Option, + pub tool_scope: ToolScope, pub declared_function_names: HashSet, @@ -367,6 +371,7 @@ impl RequestContext { rag_key: None, agent: None, last_message: None, + last_token_usage: None, tool_scope: ToolScope::default(), declared_function_names: Default::default(), node_job_scope: None, @@ -431,6 +436,7 @@ impl RequestContext { rag_key: None, agent: None, last_message: None, + last_token_usage: None, tool_scope: ToolScope { functions, mcp_runtime, @@ -492,6 +498,7 @@ impl RequestContext { rag_key: self.rag_key.clone(), agent: self.agent.clone(), last_message: self.last_message.clone(), + last_token_usage: self.last_token_usage.clone(), tool_scope: self.tool_scope.clone(), declared_function_names: self.declared_function_names.clone(), node_job_scope: None, @@ -541,6 +548,7 @@ impl RequestContext { rag_key: None, agent: None, last_message: None, + last_token_usage: None, tool_scope: ToolScope { functions: Functions::default(), mcp_runtime: McpRuntime::default(), @@ -2206,6 +2214,16 @@ impl RequestContext { } } + pub fn record_token_usage(&mut self, usage: Option) { + if let Some(usage) = &usage { + debug!("token-usage: {usage:?}"); + if let Some(session) = self.session.as_mut() { + session.accumulate_token_usage(usage); + } + } + self.last_token_usage = usage; + } + pub fn generate_prompt_context(&self, app: &AppConfig) -> HashMap<&str, String> { let mut output = HashMap::new(); let role = self.extract_role_impl(app, false).unwrap_or_else(|err| { @@ -2258,6 +2276,20 @@ impl RequestContext { if !role.is_derived() { output.insert("role", role.name().to_string()); } + if let Some(usage) = &self.last_token_usage { + if let Some(value) = usage.input_tokens { + output.insert("last_input_tokens", value.to_string()); + } + if let Some(value) = usage.output_tokens { + output.insert("last_output_tokens", value.to_string()); + } + if let Some(value) = usage.cache_read_input_tokens { + output.insert("last_cache_read_tokens", value.to_string()); + } + if let Some(value) = usage.cache_creation_input_tokens { + output.insert("last_cache_creation_tokens", value.to_string()); + } + } if let Some(session) = &self.session { output.insert("session", session.name().to_string()); if let Some(autoname) = session.autoname() { @@ -2268,6 +2300,19 @@ impl RequestContext { output.insert("consume_tokens", tokens.to_string()); output.insert("consume_percent", percent.to_string()); output.insert("user_messages_len", session.user_messages_len().to_string()); + let usage = session.token_usage(); + if let Some(value) = usage.input_tokens { + output.insert("total_input_tokens", value.to_string()); + } + if let Some(value) = usage.output_tokens { + output.insert("total_output_tokens", value.to_string()); + } + if let Some(value) = usage.cache_read_input_tokens { + output.insert("total_cache_read_tokens", value.to_string()); + } + if let Some(value) = usage.cache_creation_input_tokens { + output.insert("total_cache_creation_tokens", value.to_string()); + } } if let Some(rag) = &self.rag { output.insert("rag", rag.name().to_string()); @@ -5366,6 +5411,72 @@ mod tests { RequestContext::new(default_app_state(), WorkingMode::Cmd) } + #[test] + fn generate_prompt_context_emits_token_usage_vars() { + let mut ctx = create_test_ctx(); + ctx.session = Some(Session::default()); + ctx.record_token_usage(Some(TokenUsage { + input_tokens: Some(120), + output_tokens: Some(8), + cache_creation_input_tokens: None, + cache_read_input_tokens: Some(100), + })); + ctx.record_token_usage(Some(TokenUsage { + input_tokens: Some(30), + output_tokens: Some(2), + cache_creation_input_tokens: Some(40), + cache_read_input_tokens: None, + })); + + let app = ctx.app.config.clone(); + let vars = ctx.generate_prompt_context(&app); + + assert_eq!( + vars.get("last_input_tokens").map(String::as_str), + Some("30") + ); + assert_eq!( + vars.get("last_output_tokens").map(String::as_str), + Some("2") + ); + assert_eq!( + vars.get("last_cache_creation_tokens").map(String::as_str), + Some("40") + ); + assert!(!vars.contains_key("last_cache_read_tokens")); + assert_eq!( + vars.get("total_input_tokens").map(String::as_str), + Some("150") + ); + assert_eq!( + vars.get("total_output_tokens").map(String::as_str), + Some("10") + ); + assert_eq!( + vars.get("total_cache_read_tokens").map(String::as_str), + Some("100") + ); + assert_eq!( + vars.get("total_cache_creation_tokens").map(String::as_str), + Some("40") + ); + } + + #[test] + fn generate_prompt_context_omits_token_usage_vars_when_provider_reports_none() { + let mut ctx = create_test_ctx(); + ctx.record_token_usage(Some(TokenUsage { + input_tokens: Some(1), + ..Default::default() + })); + ctx.record_token_usage(None); + + let app = ctx.app.config.clone(); + let vars = ctx.generate_prompt_context(&app); + assert!(!vars.contains_key("last_input_tokens")); + assert!(!vars.contains_key("total_input_tokens")); + } + fn test_decl(name: &str) -> FunctionDeclaration { FunctionDeclaration { name: name.to_string(), diff --git a/src/config/session.rs b/src/config/session.rs index e28e906..fbfe6f9 100644 --- a/src/config/session.rs +++ b/src/config/session.rs @@ -1,7 +1,8 @@ use super::input::*; +use super::todo::TodoList; use super::*; -use crate::client::{Message, MessageContent, MessageRole}; +use crate::client::{Message, MessageContent, MessageRole, TokenUsage}; use crate::render::MarkdownRender; use anyhow::{Context, Result, bail}; @@ -94,6 +95,8 @@ pub struct Session { messages: Vec, #[serde(default, skip_serializing_if = "HashMap::is_empty")] data_urls: HashMap, + #[serde(default, skip_serializing_if = "TokenUsage::is_empty")] + token_usage: TokenUsage, #[serde(skip)] model: Model, @@ -233,6 +236,18 @@ impl Session { self.tokens = self.model().total_tokens(&self.messages); } + pub fn token_usage(&self) -> &TokenUsage { + &self.token_usage + } + + /// Accumulates usage reported by one API call. Each tool-loop iteration + /// is a separately billed request, so totals grow per call; compression + /// rewrites messages but never resets these counters. + pub fn accumulate_token_usage(&mut self, usage: &TokenUsage) { + self.token_usage.accumulate(usage); + self.dirty = true; + } + pub fn has_user_messages(&self) -> bool { self.messages.iter().any(|v| v.role.is_user()) } @@ -302,6 +317,18 @@ impl Session { if percent != 0.0 { data["total/max"] = format!("{percent}%").into(); } + if let Some(value) = self.token_usage.input_tokens { + data["real_input_tokens"] = value.into(); + } + if let Some(value) = self.token_usage.output_tokens { + data["real_output_tokens"] = value.into(); + } + if let Some(value) = self.token_usage.cache_read_input_tokens { + data["cache_read_tokens"] = value.into(); + } + if let Some(value) = self.token_usage.cache_creation_input_tokens { + data["cache_creation_tokens"] = value.into(); + } data["messages"] = json!(self.messages); let output = serde_yaml::to_string(&data) @@ -400,6 +427,19 @@ impl Session { items.push(("max_input_tokens", max_input_tokens.to_string())); } + if let Some(value) = self.token_usage.input_tokens { + items.push(("real_input_tokens", value.to_string())); + } + if let Some(value) = self.token_usage.output_tokens { + items.push(("real_output_tokens", value.to_string())); + } + if let Some(value) = self.token_usage.cache_read_input_tokens { + items.push(("cache_read_tokens", value.to_string())); + } + if let Some(value) = self.token_usage.cache_creation_input_tokens { + items.push(("cache_creation_tokens", value.to_string())); + } + let mut lines: Vec = items .iter() .map(|(name, value)| format!("{name:<20}{value}")) @@ -1005,7 +1045,9 @@ impl AutoName { #[cfg(test)] mod tests { use super::*; - use crate::client::{Message, MessageContent, MessageContentToolCalls, MessageRole, Model}; + use crate::client::{ + Message, MessageContent, MessageContentToolCalls, MessageRole, Model, TokenUsage, + }; use crate::config::{AppConfig, AppState, RequestContext, WorkingMode}; use crate::function::{Functions, ToolCall, ToolResult}; use std::sync::Arc; @@ -1107,6 +1149,60 @@ mod tests { } #[test] + fn session_token_usage_accumulates_and_survives_yaml_round_trip() { + let mut session = Session::default(); + assert!(session.token_usage().is_empty()); + + session.accumulate_token_usage(&TokenUsage { + input_tokens: Some(100), + output_tokens: Some(10), + cache_creation_input_tokens: Some(50), + cache_read_input_tokens: None, + }); + session.accumulate_token_usage(&TokenUsage { + input_tokens: Some(1), + output_tokens: None, + cache_creation_input_tokens: None, + cache_read_input_tokens: Some(90), + }); + assert!(session.dirty()); + assert_eq!(session.token_usage().input_tokens, Some(101)); + assert_eq!(session.token_usage().output_tokens, Some(10)); + assert_eq!(session.token_usage().cache_creation_input_tokens, Some(50)); + assert_eq!(session.token_usage().cache_read_input_tokens, Some(90)); + + let yaml = serde_yaml::to_string(&session).unwrap(); + let reloaded: Session = serde_yaml::from_str(&yaml).unwrap(); + assert_eq!(reloaded.token_usage().input_tokens, Some(101)); + assert_eq!(reloaded.token_usage().cache_read_input_tokens, Some(90)); + } + + #[test] + fn session_empty_token_usage_is_not_serialized() { + let session = Session::default(); + let yaml = serde_yaml::to_string(&session).unwrap(); + assert!(!yaml.contains("token_usage")); + } + + #[test] + fn session_export_includes_real_usage_only_when_reported() { + let empty = Session::default().export().unwrap(); + assert!(!empty.contains("real_input_tokens")); + + let mut session = Session::default(); + session.accumulate_token_usage(&TokenUsage { + input_tokens: Some(5), + output_tokens: Some(7), + cache_creation_input_tokens: None, + cache_read_input_tokens: Some(3), + }); + let exported = session.export().unwrap(); + assert!(exported.contains("real_input_tokens: 5")); + assert!(exported.contains("real_output_tokens: 7")); + assert!(exported.contains("cache_read_tokens: 3")); + assert!(!exported.contains("cache_creation_tokens")); + } + fn session_mcp_tools_survives_yaml_round_trip() { let mut session = Session::default(); let mut mcp_tools = IndexMap::new();