From e44efa709e3ea12bfbcdd3f83a3879164491ce6b Mon Sep 17 00:00:00 2001 From: whit3rabbit Date: Sat, 12 Sep 2026 08:59:52 -0500 Subject: [PATCH] Account for Anthropic web search costs --- crates/proxy/src/cost/mod.rs | 75 +++++++++++++++++-- crates/proxy/src/server/chat_completions.rs | 18 ++++- .../src/server/chat_completions/stream.rs | 13 +++- crates/proxy/src/server/routes.rs | 20 ++++- crates/proxy/src/server/streaming.rs | 29 +++++++ 5 files changed, 142 insertions(+), 13 deletions(-) diff --git a/crates/proxy/src/cost/mod.rs b/crates/proxy/src/cost/mod.rs index efc7f61..335049c 100644 --- a/crates/proxy/src/cost/mod.rs +++ b/crates/proxy/src/cost/mod.rs @@ -214,6 +214,26 @@ impl ModelPricing { } } +const ANTHROPIC_WEB_SEARCH_COST_USD: f64 = 0.01; + +/// Calculate the non-token cost Anthropic reports for hosted server tools. +/// +/// Anthropic bills web search at $10 per 1,000 successful searches, in addition +/// to the model tokens reported in the same usage object. Other server-tool +/// counters are deliberately ignored until they have separate billable pricing. +pub fn anthropic_web_search_cost(requests: u64) -> f64 { + requests as f64 * ANTHROPIC_WEB_SEARCH_COST_USD +} + +pub fn anthropic_server_tool_usage_cost( + usage: Option<&anyllm_translate::anthropic::messages::ServerToolUsage>, +) -> f64 { + usage + .and_then(|usage| usage.web_search_requests) + .map(|requests| anthropic_web_search_cost(requests as u64)) + .unwrap_or(0.0) +} + /// Record cost for a completed request against a virtual key. /// /// Calculates cost from token usage and the resolved model name, then @@ -226,7 +246,24 @@ pub fn record_cost( input_tokens: u64, output_tokens: u64, ) -> f64 { - let cost = pricing().cost_for_usage(model, input_tokens, output_tokens); + record_cost_with_extra(shared, vk_ctx, model, input_tokens, output_tokens, 0.0) +} + +/// Record token cost plus an explicit non-token surcharge against a virtual key. +pub fn record_cost_with_extra( + shared: &Option, + vk_ctx: &Option, + model: &str, + input_tokens: u64, + output_tokens: u64, + extra_cost_usd: f64, +) -> f64 { + let token_cost = if input_tokens == 0 && output_tokens == 0 { + 0.0 + } else { + pricing().cost_for_usage(model, input_tokens, output_tokens) + }; + let cost = token_cost + extra_cost_usd.max(0.0); if cost <= 0.0 { return cost; } @@ -295,6 +332,19 @@ mod tests { } } + #[test] + fn anthropic_server_tool_usage_cost_counts_web_search_only() { + let usage = anyllm_translate::anthropic::messages::ServerToolUsage { + web_search_requests: Some(3), + web_fetch_requests: Some(9), + tool_search_requests: Some(7), + extra: serde_json::Map::new(), + }; + + assert!((anthropic_server_tool_usage_cost(Some(&usage)) - 0.03).abs() < 1e-12); + assert_eq!(anthropic_server_tool_usage_cost(None), 0.0); + } + #[test] fn exact_match() { let pricing = test_pricing(); @@ -467,24 +517,39 @@ mod tests { period_reset: None, }; + let shared_opt = Some(shared); + let vk_ctx_opt = Some(vk_ctx); + // record_cost uses tokio::task::spawn_blocking, so we need a runtime. let rt = tokio::runtime::Runtime::new().unwrap(); rt.block_on(async { - let cost = record_cost(&Some(shared), &Some(vk_ctx), "gpt-4o", 1000, 500); + let cost = record_cost(&shared_opt, &vk_ctx_opt, "gpt-4o", 1000, 500); assert!(cost > 0.0); // Wait for the spawned blocking task to complete. tokio::task::yield_now().await; tokio::time::sleep(std::time::Duration::from_millis(100)).await; + + let extra_only_cost = record_cost_with_extra( + &shared_opt, + &vk_ctx_opt, + "unknown-anthropic-model", + 0, + 0, + 0.03, + ); + assert!((extra_only_cost - 0.03).abs() < 1e-12); + tokio::task::yield_now().await; + tokio::time::sleep(std::time::Duration::from_millis(100)).await; }); - // Verify the spend was persisted. + // Verify the spend was persisted, including an extra-cost-only hosted tool charge. let conn = db.lock().unwrap(); let spend = db::get_key_spend(&conn, key_id).unwrap().unwrap(); - assert!(spend.total_cost_usd > 0.0); + assert!(spend.total_cost_usd > 0.03); assert_eq!(spend.total_input_tokens, 1000); assert_eq!(spend.total_output_tokens, 500); - assert_eq!(spend.request_count, 1); + assert_eq!(spend.request_count, 2); } // -- Spend threshold detection tests -- diff --git a/crates/proxy/src/server/chat_completions.rs b/crates/proxy/src/server/chat_completions.rs index fed78dd..d585916 100644 --- a/crates/proxy/src/server/chat_completions.rs +++ b/crates/proxy/src/server/chat_completions.rs @@ -333,12 +333,16 @@ pub(crate) async fn chat_completions( let oai_response = translate_anthropic_to_openai_response(&anthropic_resp, &original_model); - let cost = super::routes::record_virtual_key_usage( + let server_tool_cost = crate::cost::anthropic_server_tool_usage_cost( + anthropic_resp.usage.server_tool_use.as_ref(), + ); + let cost = super::routes::record_virtual_key_usage_with_extra( &state.shared, &vk_ctx, &mapped_model, anthropic_resp.usage.input_tokens as u64, anthropic_resp.usage.output_tokens as u64, + server_tool_cost, ); log_request( &state.shared, @@ -415,12 +419,16 @@ pub(crate) async fn chat_completions( ); let oai_response = translate_anthropic_to_openai_response(&anthropic_resp, &original_model); - let cost = super::routes::record_virtual_key_usage( + let server_tool_cost = crate::cost::anthropic_server_tool_usage_cost( + anthropic_resp.usage.server_tool_use.as_ref(), + ); + let cost = super::routes::record_virtual_key_usage_with_extra( &state.shared, &vk_ctx, &mapped_model, anthropic_resp.usage.input_tokens as u64, anthropic_resp.usage.output_tokens as u64, + server_tool_cost, ); log_request( &state.shared, @@ -536,12 +544,16 @@ pub(crate) async fn chat_completions( &original_model, &tool_context, ); - let cost = super::routes::record_virtual_key_usage( + let server_tool_cost = crate::cost::anthropic_server_tool_usage_cost( + anthropic_resp.usage.server_tool_use.as_ref(), + ); + let cost = super::routes::record_virtual_key_usage_with_extra( &state.shared, &vk_ctx, &mapped_model, anthropic_resp.usage.input_tokens as u64, anthropic_resp.usage.output_tokens as u64, + server_tool_cost, ); log_request( &state.shared, diff --git a/crates/proxy/src/server/chat_completions/stream.rs b/crates/proxy/src/server/chat_completions/stream.rs index 6961fd6..1be33a8 100644 --- a/crates/proxy/src/server/chat_completions/stream.rs +++ b/crates/proxy/src/server/chat_completions/stream.rs @@ -199,15 +199,20 @@ async fn anthropic_chat_completions_stream( } let tokens = usage.tokens(); - let cost = tokens.map(|(input_t, output_t)| { - crate::server::routes::record_virtual_key_usage( + let server_tool_cost = crate::cost::anthropic_web_search_cost(usage.web_search_requests()); + let cost = if tokens.is_some() || server_tool_cost > 0.0 { + let (input_t, output_t) = tokens.unwrap_or((0, 0)); + Some(crate::server::routes::record_virtual_key_usage_with_extra( &log_shared, &vk_ctx, &mapped_model, input_t, output_t, - ) - }); + server_tool_cost, + )) + } else { + None + }; let (status, err) = outcome.record(&metrics); log_request( &log_shared, diff --git a/crates/proxy/src/server/routes.rs b/crates/proxy/src/server/routes.rs index 8588b6f..90668f2 100644 --- a/crates/proxy/src/server/routes.rs +++ b/crates/proxy/src/server/routes.rs @@ -1291,10 +1291,28 @@ pub(crate) fn record_virtual_key_usage( model: &str, input_tokens: u64, output_tokens: u64, +) -> f64 { + record_virtual_key_usage_with_extra(shared, vk_ctx, model, input_tokens, output_tokens, 0.0) +} + +pub(crate) fn record_virtual_key_usage_with_extra( + shared: &Option, + vk_ctx: &Option, + model: &str, + input_tokens: u64, + output_tokens: u64, + extra_cost_usd: f64, ) -> f64 { let capped_output = output_tokens.min(u32::MAX as u64) as u32; record_vk_tpm(vk_ctx, capped_output); - crate::cost::record_cost(shared, vk_ctx, model, input_tokens, output_tokens) + crate::cost::record_cost_with_extra( + shared, + vk_ctx, + model, + input_tokens, + output_tokens, + extra_cost_usd, + ) } /// Global webhook callback config, set once at startup. diff --git a/crates/proxy/src/server/streaming.rs b/crates/proxy/src/server/streaming.rs index d8c40cb..dded08f 100644 --- a/crates/proxy/src/server/streaming.rs +++ b/crates/proxy/src/server/streaming.rs @@ -82,12 +82,23 @@ data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop assert_eq!(usage.tokens(), Some((11, 7))); assert!(buffer.is_empty()); } + + #[test] + fn anthropic_stream_usage_counts_successful_web_search_results() { + let mut usage = AnthropicStreamUsage::default(); + + usage.observe_data(r#"{"type":"content_block_start","index":1,"content_block":{"type":"web_search_tool_result","tool_use_id":"srvtoolu_ok","content":[]}}"#); + usage.observe_data(r#"{"type":"content_block_start","index":2,"content_block":{"type":"web_search_tool_result","tool_use_id":"srvtoolu_err","is_error":true}}"#); + + assert_eq!(usage.web_search_requests(), 1); + } } #[derive(Debug, Default, Clone, Copy)] pub(crate) struct AnthropicStreamUsage { input_tokens: Option, output_tokens: Option, + web_search_requests: u64, } impl AnthropicStreamUsage { @@ -107,6 +118,20 @@ impl AnthropicStreamUsage { } => { self.output_tokens = Some(usage.output_tokens as u64); } + anthropic::StreamEvent::ContentBlockStart { + content_block: + anthropic::ContentBlock::WebSearchToolResult { + is_error: Some(true), + .. + }, + .. + } => {} + anthropic::StreamEvent::ContentBlockStart { + content_block: anthropic::ContentBlock::WebSearchToolResult { .. }, + .. + } => { + self.web_search_requests = self.web_search_requests.saturating_add(1); + } _ => {} } } @@ -119,6 +144,10 @@ impl AnthropicStreamUsage { (None, None) => None, } } + + pub(crate) fn web_search_requests(&self) -> u64 { + self.web_search_requests + } } pub(crate) fn observe_anthropic_sse_frames(