diff --git a/crates/goose-provider-types/src/formats/openai_responses.rs b/crates/goose-provider-types/src/formats/openai_responses.rs index 71a06090a1..a172558a29 100644 --- a/crates/goose-provider-types/src/formats/openai_responses.rs +++ b/crates/goose-provider-types/src/formats/openai_responses.rs @@ -150,21 +150,22 @@ pub struct ResponseUsage { pub struct InputTokensDetails { #[serde(default)] pub cached_tokens: Option, + #[serde(default)] + pub cache_write_tokens: Option, } impl ResponseUsage { fn to_usage(&self) -> Usage { - // input_tokens already includes cached tokens - let cached_tokens = self - .input_tokens_details - .as_ref() - .and_then(|d| d.cached_tokens); + // input_tokens already includes both cache reads and cache writes + let details = self.input_tokens_details.as_ref(); + let cached_tokens = details.and_then(|d| d.cached_tokens); + let cache_write_tokens = details.and_then(|d| d.cache_write_tokens); Usage::new( Some(self.input_tokens), Some(self.output_tokens), Some(self.total_tokens), ) - .with_cache_tokens(cached_tokens, None) + .with_cache_tokens(cached_tokens, cache_write_tokens) } } diff --git a/crates/goose-provider-types/tests/openai_responses_cache_usage.rs b/crates/goose-provider-types/tests/openai_responses_cache_usage.rs new file mode 100644 index 0000000000..2c818cdf65 --- /dev/null +++ b/crates/goose-provider-types/tests/openai_responses_cache_usage.rs @@ -0,0 +1,76 @@ +use futures::StreamExt; +use goose_provider_types::conversation::token_usage::Usage; +use goose_provider_types::formats::openai_responses::{ + get_responses_usage, responses_api_to_streaming_message, ResponsesApiResponse, +}; +use serde_json::json; + +async fn final_stream_usage(lines: Vec) -> Usage { + let messages = + responses_api_to_streaming_message(tokio_stream::iter(lines.into_iter().map(Ok))); + futures::pin_mut!(messages); + + let mut usage = None; + while let Some(item) = messages.next().await { + let (_, maybe_usage) = item.expect("stream item should parse"); + if maybe_usage.is_some() { + usage = maybe_usage; + } + } + usage.expect("stream should report usage").usage +} + +#[tokio::test] +async fn stream_completed_reports_cache_write_tokens() { + let usage = final_stream_usage(vec![ + r#"data: {"type":"response.created","sequence_number":1,"response":{"id":"resp_1","object":"response","created_at":1737368310,"status":"in_progress","model":"gpt-6-luna","output":[]}}"#.to_string(), + r#"data: {"type":"response.completed","sequence_number":2,"response":{"id":"resp_1","object":"response","created_at":1737368310,"status":"completed","model":"gpt-6-luna","output":[],"usage":{"input_tokens":5457,"input_tokens_details":{"cache_write_tokens":5454,"cached_tokens":0},"output_tokens":5,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":5462}}}"#.to_string(), + "data: [DONE]".to_string(), + ]) + .await; + + assert_eq!(usage.input_tokens, Some(5457)); + assert_eq!(usage.total_tokens, Some(5462)); + assert_eq!(usage.cache_read_input_tokens, Some(0)); + assert_eq!(usage.cache_write_input_tokens, Some(5454)); +} + +#[tokio::test] +async fn stream_incomplete_reports_cache_write_tokens() { + let usage = final_stream_usage(vec![ + r#"data: {"type":"response.created","sequence_number":1,"response":{"id":"resp_1","object":"response","created_at":1737368310,"status":"in_progress","model":"gpt-6-luna","output":[]}}"#.to_string(), + r#"data: {"type":"response.incomplete","sequence_number":2,"response":{"id":"resp_1","object":"response","created_at":1737368310,"status":"incomplete","model":"gpt-6-luna","output":[],"incomplete_details":{"reason":"max_output_tokens"},"usage":{"input_tokens":10,"input_tokens_details":{"cache_write_tokens":6,"cached_tokens":3},"output_tokens":5,"total_tokens":15}}}"#.to_string(), + "data: [DONE]".to_string(), + ]) + .await; + + assert_eq!(usage.input_tokens, Some(10)); + assert_eq!(usage.cache_read_input_tokens, Some(3)); + assert_eq!(usage.cache_write_input_tokens, Some(6)); +} + +#[test] +fn non_streaming_usage_reports_cache_read_and_write_tokens() { + let response: ResponsesApiResponse = serde_json::from_value(json!({ + "id": "resp_1", + "object": "response", + "created_at": 1737368310, + "status": "completed", + "model": "gpt-6-luna", + "output": [], + "usage": { + "input_tokens": 2176, + "input_tokens_details": { "cache_write_tokens": 9, "cached_tokens": 2164 }, + "output_tokens": 7, + "output_tokens_details": { "reasoning_tokens": 0 }, + "total_tokens": 2183 + } + })) + .unwrap(); + + let usage = get_responses_usage(&response); + + assert_eq!(usage.input_tokens, Some(2176)); + assert_eq!(usage.cache_read_input_tokens, Some(2164)); + assert_eq!(usage.cache_write_input_tokens, Some(9)); +}