feat: Support token caching to reduce LLM costs in Anthropic models (not typically handled automatically by APIs)
This commit is contained in:
10 files changed
+984
-78
No files matched your search
+2
-1
@@ -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<Option<String>> = OnceLock::new();
|
||||
|
||||
|
||||
@@ -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<LastMessage>,
|
||||
|
||||
/// 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<TokenUsage>,
|
||||
|
||||
pub tool_scope: ToolScope,
|
||||
|
||||
pub declared_function_names: HashSet<String>,
|
||||
@@ -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<TokenUsage>) {
|
||||
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(),
|
||||
|
||||
+98
-2
@@ -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<Message>,
|
||||
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
|
||||
data_urls: HashMap<String, String>,
|
||||
#[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<String> = 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();
|
||||
|
||||
Reference in new issue
Block a user