diff --git a/crates/anthropic/src/anthropic.rs b/crates/anthropic/src/anthropic.rs index 178cd8d831..ed1e761c89 100644 --- a/crates/anthropic/src/anthropic.rs +++ b/crates/anthropic/src/anthropic.rs @@ -423,6 +423,45 @@ pub async fn non_streaming_completion( } } +/// Estimates input tokens without generating a message. +/// +/// Anthropic's estimate may differ slightly from usage reported during generation. +pub async fn count_input_tokens( + client: &dyn HttpClient, + api_url: &str, + api_key: &str, + request: CountTokensRequest, + beta_headers: Option, + extra_headers: &CustomHeaders, +) -> Result { + let (mut response, rate_limits) = send_request_to_route( + client, + api_url, + "/v1/messages/count_tokens", + api_key, + request, + beta_headers, + extra_headers, + ) + .await?; + if !response.status().is_success() { + return Err(handle_error_response(response, rate_limits).await); + } + let mut body = String::new(); + response + .body_mut() + .read_to_string(&mut body) + .await + .map_err(AnthropicError::ReadResponse)?; + #[derive(Deserialize)] + struct CountTokensResponse { + input_tokens: u64, + } + serde_json::from_str::(&body) + .map(|response| response.input_tokens) + .map_err(AnthropicError::DeserializeResponse) +} + async fn send_request( client: &dyn HttpClient, api_url: &str, @@ -431,7 +470,28 @@ async fn send_request( beta_headers: Option, extra_headers: &CustomHeaders, ) -> Result<(http::Response, RateLimitInfo), AnthropicError> { - let uri = format!("{api_url}/v1/messages"); + send_request_to_route( + client, + api_url, + "/v1/messages", + api_key, + request, + beta_headers, + extra_headers, + ) + .await +} + +async fn send_request_to_route( + client: &dyn HttpClient, + api_url: &str, + route: &str, + api_key: &str, + request: impl Serialize, + beta_headers: Option, + extra_headers: &CustomHeaders, +) -> Result<(http::Response, RateLimitInfo), AnthropicError> { + let uri = format!("{api_url}{route}"); let mut request_builder = HttpRequest::builder() .method(Method::POST) @@ -949,7 +1009,36 @@ pub struct Request { pub top_p: Option, } +/// Input counted by Anthropic, excluding generation and service settings. +#[derive(Serialize, Debug)] +pub struct CountTokensRequest { + pub model: String, + pub messages: Vec, + #[serde(skip_serializing_if = "Vec::is_empty")] + pub tools: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_choice: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub system: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub thinking: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub context_management: Option, +} + impl Request { + pub fn into_count_tokens_request(self) -> CountTokensRequest { + CountTokensRequest { + model: self.model, + messages: self.messages, + tools: self.tools, + tool_choice: self.tool_choice, + system: self.system, + thinking: self.thinking, + context_management: self.context_management, + } + } + /// Configures this request to stop after native compaction. /// /// Tool definitions remain in the request so the trigger observes the same @@ -1357,6 +1446,115 @@ mod tests { use super::*; use http_client::FakeHttpClient; + #[test] + fn count_input_tokens_preserves_input_without_generation_fields() { + let request: Request = serde_json::from_value(serde_json::json!({ + "model": "claude-sonnet-4-6", + "max_tokens": 1234, + "temperature": 0.5, + "system": "Be precise", + "messages": [{"role": "user", "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image", "source": { + "type": "base64", "media_type": "image/png", "data": "aW1hZ2U=" + }} + ]}], + "tools": [{"name": "look", "description": "Look up data", "input_schema": {"type": "object"}}], + "tool_choice": {"type": "auto"}, + "context_management": {"edits": [{"type": "compact_20260112"}]} + })).expect("generation request"); + let generation = serde_json::to_value(&request).expect("generation payload"); + assert_eq!(generation["max_tokens"], 1234); + assert_eq!(generation["temperature"], 0.5); + let client = FakeHttpClient::create(move |mut request| async move { + assert_eq!(request.method(), Method::POST); + assert_eq!( + request.uri(), + "https://api.anthropic.com/v1/messages/count_tokens" + ); + assert!( + request + .headers() + .get("X-Api-Key") + .is_some_and(|value| value == "test-key") + ); + assert_eq!(request.headers()["Anthropic-Version"], "2023-06-01"); + assert_eq!(request.headers()["Anthropic-Beta"], "test-beta"); + let mut body = String::new(); + request.body_mut().read_to_string(&mut body).await?; + assert_eq!( + serde_json::from_str::(&body)?, + serde_json::json!({ + "model": "claude-sonnet-4-6", + "system": "Be precise", + "messages": [{"role": "user", "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image", "source": { + "type": "base64", "media_type": "image/png", "data": "aW1hZ2U=" + }} + ]}], + "tools": [{"name": "look", "description": "Look up data", "input_schema": {"type": "object"}}], + "tool_choice": {"type": "auto"}, + "context_management": {"edits": [{"type": "compact_20260112"}]} + }) + ); + Ok(http::Response::builder() + .status(200) + .body(AsyncBody::from(r#"{"input_tokens":321}"#))?) + }); + let count = futures::executor::block_on(count_input_tokens( + client.as_ref(), + ANTHROPIC_API_URL, + " test-key ", + request.into_count_tokens_request(), + Some("test-beta".into()), + &CustomHeaders::default(), + )) + .expect("count succeeds"); + assert_eq!(count, 321); + } + + #[test] + fn count_input_tokens_preserves_typed_errors() { + for (status, body) in [ + ( + 401, + r#"{"type":"error","error":{"type":"authentication_error","message":"invalid x-api-key"}}"#, + ), + (200, "{}"), + ] { + let client = FakeHttpClient::create(move |_| async move { + Ok(http::Response::builder() + .status(status) + .body(AsyncBody::from(body))?) + }); + let request: Request = serde_json::from_value(serde_json::json!({ + "model": "claude-sonnet-4-6", "max_tokens": 100, "messages": [] + })) + .expect("request"); + let error = futures::executor::block_on(count_input_tokens( + client.as_ref(), + ANTHROPIC_API_URL, + "test-key", + request.into_count_tokens_request(), + None, + &CustomHeaders::default(), + )) + .expect_err("count fails"); + if status == 401 { + assert!(matches!( + error, + AnthropicError::ApiError { + status: Some(StatusCode::UNAUTHORIZED), + .. + } + )); + } else { + assert!(matches!(error, AnthropicError::DeserializeResponse(_))); + } + } + } + #[test] fn list_models_preserves_anthropic_api_errors() { let client = FakeHttpClient::create(|_| async move { diff --git a/crates/language_model/src/fake_provider.rs b/crates/language_model/src/fake_provider.rs index a9cb4c5b5f..94f4487f65 100644 --- a/crates/language_model/src/fake_provider.rs +++ b/crates/language_model/src/fake_provider.rs @@ -9,6 +9,7 @@ use futures::{FutureExt, channel::mpsc, future::BoxFuture, stream::BoxStream, st use gpui::{App, AsyncApp, Entity, Task}; use http_client::Result; use parking_lot::Mutex; +use std::collections::VecDeque; use std::sync::{ Arc, atomic::{AtomicBool, AtomicU64, Ordering::SeqCst}, @@ -121,6 +122,8 @@ pub struct FakeLanguageModel { supports_server_side_compaction: AtomicBool, max_token_count: AtomicU64, max_output_tokens: AtomicU64, + input_token_counts: Mutex>, + input_token_count_requests: Mutex>, } impl Default for FakeLanguageModel { @@ -139,6 +142,8 @@ impl Default for FakeLanguageModel { supports_server_side_compaction: AtomicBool::new(false), max_token_count: AtomicU64::new(1_000_000), max_output_tokens: AtomicU64::new(0), + input_token_counts: Mutex::new(VecDeque::new()), + input_token_count_requests: Mutex::new(Vec::new()), } } } @@ -196,6 +201,14 @@ impl FakeLanguageModel { .store(count.unwrap_or_default(), SeqCst); } + pub fn queue_input_token_count(&self, count: u64) { + self.input_token_counts.lock().push_back(count); + } + + pub fn input_token_count_requests(&self) -> Vec { + self.input_token_count_requests.lock().clone() + } + pub fn pending_completions(&self) -> Vec { self.current_completion_txs .lock() @@ -338,6 +351,24 @@ impl LanguageModel for FakeLanguageModel { } } + fn count_input_tokens( + &self, + request: LanguageModelRequest, + _: &AsyncApp, + ) -> BoxFuture<'static, Result, LanguageModelCompletionError>> { + if self.forbid_requests.load(SeqCst) { + return async { + Err(LanguageModelCompletionError::Other(anyhow!( + "requests are forbidden" + ))) + } + .boxed(); + } + self.input_token_count_requests.lock().push(request); + let count = self.input_token_counts.lock().pop_front(); + async move { Ok(count) }.boxed() + } + fn stream_completion( &self, request: LanguageModelRequest, diff --git a/crates/language_model/src/language_model.rs b/crates/language_model/src/language_model.rs index 38c89fd3e1..90353facb8 100644 --- a/crates/language_model/src/language_model.rs +++ b/crates/language_model/src/language_model.rs @@ -218,6 +218,19 @@ pub trait LanguageModel: Send + Sync { /// Returns the input token ceiling. fn max_token_count(&self) -> u64; + /// Counts request input without generating output, when supported by the provider. + /// + /// Counts may be estimates and differ from subsequent measured usage. Callers + /// choose the content to count; this does not infer which input is already + /// covered by a previous usage report. Unsupported providers return `None`. + fn count_input_tokens( + &self, + _request: LanguageModelRequest, + _cx: &AsyncApp, + ) -> BoxFuture<'static, Result, LanguageModelCompletionError>> { + async { Ok(None) }.boxed() + } + /// Returns the combined input and output ceiling, if one applies. /// /// The conservative default shares the input ceiling with output. `None` diff --git a/crates/language_models/src/provider/anthropic.rs b/crates/language_models/src/provider/anthropic.rs index 686ff2b702..e8d380b00f 100644 --- a/crates/language_models/src/provider/anthropic.rs +++ b/crates/language_models/src/provider/anthropic.rs @@ -892,6 +892,51 @@ impl LanguageModel for AnthropicModel { Some(self.model.max_output_tokens) } + fn count_input_tokens( + &self, + request: LanguageModelRequest, + cx: &AsyncApp, + ) -> BoxFuture<'static, Result, LanguageModelCompletionError>> { + let request_id = self.model.request_id(!request.tools.is_empty()).to_string(); + let request = into_anthropic( + request, + request_id, + self.model.default_temperature, + self.model.max_output_tokens, + self.model.mode.clone(), + AnthropicPromptCacheMode::Automatic, + &PROVIDER_ID, + ); + let http_client = self.http_client.clone(); + let (api_key, api_url, extra_headers) = self.state.read_with(cx, |state, cx| { + let api_url = AnthropicLanguageModelProvider::api_url(cx); + let extra_headers = AnthropicLanguageModelProvider::settings(cx) + .custom_headers + .clone(); + (state.api_key_state.key(&api_url), api_url, extra_headers) + }); + let beta_headers = self.model.beta_headers(); + self.request_limiter + .run(async move { + let request = request?.into_count_tokens_request(); + let api_key = api_key.ok_or(LanguageModelCompletionError::NoApiKey { + provider: PROVIDER_NAME, + })?; + anthropic::count_input_tokens( + http_client.as_ref(), + &api_url, + &api_key, + request, + beta_headers, + &extra_headers, + ) + .await + .map(Some) + .map_err(Into::into) + }) + .boxed() + } + fn stream_completion( &self, request: LanguageModelRequest, diff --git a/crates/language_models/src/provider/open_ai.rs b/crates/language_models/src/provider/open_ai.rs index 328f8005c7..009fa22f80 100644 --- a/crates/language_models/src/provider/open_ai.rs +++ b/crates/language_models/src/provider/open_ai.rs @@ -609,6 +609,56 @@ impl LanguageModel for OpenAiLanguageModel { self.model.max_output_tokens() } + fn count_input_tokens( + &self, + mut request: LanguageModelRequest, + cx: &AsyncApp, + ) -> BoxFuture<'static, Result, LanguageModelCompletionError>> { + if !self.model.uses_responses_api() { + return async { Ok(None) }.boxed(); + } + normalize_open_ai_response_thinking_effort(&mut request, &self.model); + let request = into_open_ai_response( + request, + self.model.id(), + self.model.supports_parallel_tool_calls(), + self.model.supports_prompt_cache_key(), + self.max_output_tokens(), + default_thinking_reasoning_effort(&self.model), + self.model + .supported_reasoning_efforts() + .contains(&open_ai::ReasoningEffort::None), + &OPEN_AI_PROVIDER_ID, + ); + let http_client = self.http_client.clone(); + let (api_key, api_url, extra_headers) = self.state.read_with(cx, |state, cx| { + let api_url = OpenAiLanguageModelProvider::api_url(cx); + let extra_headers = OpenAiLanguageModelProvider::settings(cx) + .custom_headers + .clone(); + (state.api_key_state.key(&api_url), api_url, extra_headers) + }); + self.request_limiter + .run(async move { + let request = request?.into_count_tokens_request(); + let api_key = api_key.ok_or(LanguageModelCompletionError::NoApiKey { + provider: PROVIDER_NAME, + })?; + open_ai::responses::count_input_tokens( + http_client.as_ref(), + PROVIDER_NAME.0.as_str(), + &api_url, + &api_key, + request, + &extra_headers, + ) + .await + .map(Some) + .map_err(Into::into) + }) + .boxed() + } + fn stream_completion( &self, mut request: LanguageModelRequest, diff --git a/crates/language_models_cloud/src/language_models_cloud.rs b/crates/language_models_cloud/src/language_models_cloud.rs index 13b899317c..6e9ff077fb 100644 --- a/crates/language_models_cloud/src/language_models_cloud.rs +++ b/crates/language_models_cloud/src/language_models_cloud.rs @@ -248,6 +248,87 @@ impl CloudLanguageModel { } impl CloudLanguageModel { + fn check_data_retention_consent( + &self, + cx: &AsyncApp, + ) -> Result<(), LanguageModelCompletionError> { + if self.requires_data_retention() && !self.token_provider.has_data_retention_consent(cx) { + return Err(LanguageModelCompletionError::DataRetentionConsentRequired { + model_name: self.model.display_name.clone(), + }); + } + Ok(()) + } + + fn anthropic_request(&self, request: LanguageModelRequest) -> Result { + let enable_thinking = request.thinking_allowed && self.model.supports_thinking; + let effort = request + .thinking_effort + .as_ref() + .and_then(|effort| anthropic::Effort::from_str(effort).ok()); + let mut request = into_anthropic( + request, + self.model.id.to_string(), + 1.0, + self.model.max_output_tokens as u64, + if enable_thinking { + AnthropicModelMode::Thinking { + budget_tokens: Some(4_096), + } + } else { + AnthropicModelMode::Default + }, + AnthropicPromptCacheMode::Automatic, + &ANTHROPIC_PROVIDER_ID, + )?; + if enable_thinking && effort.is_some() { + request.thinking = Some(anthropic::Thinking::Adaptive { + display: Some(anthropic::AdaptiveThinkingDisplay::Summarized), + // The cloud proxy owns the beta header needed for block binding. + block_binding: None, + }); + request.output_config = Some(anthropic::OutputConfig { effort }); + } + if !self.model.supports_fast_mode { + request.speed = None; + } + Ok(request) + } + + fn open_ai_request( + &self, + request: LanguageModelRequest, + ) -> Result { + let enable_thinking = request.thinking_allowed && self.model.supports_thinking; + let effort = request + .thinking_effort + .as_ref() + .and_then(|effort| open_ai::ReasoningEffort::from_str(effort).ok()) + .filter(|effort| *effort != open_ai::ReasoningEffort::None); + let supports_none_reasoning_effort = + self.model.supported_effort_levels.iter().any(|effort| { + open_ai::ReasoningEffort::from_str(&effort.value) + .is_ok_and(|effort| effort == open_ai::ReasoningEffort::None) + }); + let mut request = into_open_ai_response( + request, + &self.model.id.0, + self.model.supports_parallel_tool_calls, + true, + None, + None, + supports_none_reasoning_effort, + &OPEN_AI_PROVIDER_ID, + )?; + if enable_thinking && let Some(effort) = effort { + request.reasoning = Some(open_ai::responses::ReasoningConfig { + effort, + summary: Some(open_ai::responses::ReasoningSummaryMode::Auto), + }); + } + Ok(request) + } + fn compact_anthropic( &self, request: LanguageModelRequest, @@ -625,6 +706,9 @@ impl LanguageModel for CloudLanguageModel BoxFuture<'static, Result> { + if let Err(error) = self.check_data_retention_consent(cx) { + return async move { Err(error) }.boxed(); + } if !self.supports_explicit_compaction() { return async { Err(LanguageModelCompletionError::Other(anyhow::anyhow!( @@ -696,6 +780,75 @@ impl LanguageModel for CloudLanguageModel BoxFuture<'static, Result, LanguageModelCompletionError>> { + if let Err(error) = self.check_data_retention_consent(cx) { + return async move { Err(error) }.boxed(); + } + use cloud_llm_client::LanguageModelProvider; + let provider_request = match self.model.provider { + LanguageModelProvider::Anthropic => self + .anthropic_request(request) + .and_then(|request| Ok(serde_json::to_value(request)?)), + LanguageModelProvider::OpenAi => self + .open_ai_request(request) + .and_then(|request| Ok(serde_json::to_value(request)?)), + _ => return async { Ok(None) }.boxed(), + }; + let body = match provider_request { + Ok(provider_request) => CompletionBody { + thread_id: None, + prompt_id: None, + provider: self.model.provider, + model: self.model.id.to_string(), + provider_request, + }, + Err(error) => return async move { Err(error.into()) }.boxed(), + }; + let http_client = self.http_client.clone(); + let token_provider = self.token_provider.clone(); + let auth_context = token_provider.auth_context(cx); + let app_version = self.app_version.clone(); + self.request_limiter + .run(async move { + let mut response = Self::perform_llm_request( + "/count_tokens", + false, + &http_client, + &*token_provider, + auth_context, + app_version, + body, + ) + .await? + .response; + let mut body = String::new(); + response + .body_mut() + .read_to_string(&mut body) + .await + .map_err(|error| LanguageModelCompletionError::ApiReadResponseError { + provider: PROVIDER_NAME, + error, + })?; + #[derive(Deserialize)] + struct CountResponse { + tokens: u64, + } + let count: CountResponse = serde_json::from_str(&body).map_err(|error| { + LanguageModelCompletionError::DeserializeResponse { + provider: PROVIDER_NAME, + error, + } + })?; + Ok(Some(count.tokens)) + }) + .boxed() + } + fn stream_completion( &self, request: LanguageModelRequest, @@ -707,12 +860,8 @@ impl LanguageModel for CloudLanguageModel, > { - if self.requires_data_retention() && !self.token_provider.has_data_retention_consent(cx) { - let model_name = self.model.display_name.clone(); - return async move { - Err(LanguageModelCompletionError::DataRetentionConsentRequired { model_name }) - } - .boxed(); + if let Err(error) = self.check_data_retention_consent(cx) { + return async move { Err(error) }.boxed(); } let mut request = request; @@ -723,53 +872,14 @@ impl LanguageModel for CloudLanguageModel { - let effort = request - .thinking_effort - .as_ref() - .and_then(|effort| anthropic::Effort::from_str(effort).ok()); - - let mut request = match into_anthropic( - request, - self.model.id.to_string(), - 1.0, - self.model.max_output_tokens as u64, - if enable_thinking { - AnthropicModelMode::Thinking { - budget_tokens: Some(4_096), - } - } else { - AnthropicModelMode::Default - }, - AnthropicPromptCacheMode::Automatic, - // Cloud proxies to Anthropic's own infrastructure, so - // compaction state is owned by (and interchangeable with) - // Anthropic proper, not by the cloud transport. - &ANTHROPIC_PROVIDER_ID, - ) { + let request = match self.anthropic_request(request) { Ok(request) => request, Err(error) => return async move { Err(error.into()) }.boxed(), }; - if enable_thinking && effort.is_some() { - request.thinking = Some(anthropic::Thinking::Adaptive { - display: Some(anthropic::AdaptiveThinkingDisplay::Summarized), - // Thinking block binding needs a beta header on the - // upstream Anthropic request, which the cloud proxy - // controls, so opting in belongs server-side. - block_binding: None, - }); - request.output_config = Some(anthropic::OutputConfig { effort }); - } - - if !self.model.supports_fast_mode { - request.speed = None; - } - let http_client = self.http_client.clone(); let token_provider = self.token_provider.clone(); let auth_context = token_provider.auth_context(cx); @@ -815,38 +925,11 @@ impl LanguageModel for CloudLanguageModel { let http_client = self.http_client.clone(); let token_provider = self.token_provider.clone(); - let effort = request - .thinking_effort - .as_ref() - .and_then(|effort| open_ai::ReasoningEffort::from_str(effort).ok()) - .filter(|effort| *effort != open_ai::ReasoningEffort::None); - let supports_none_reasoning_effort = - self.model.supported_effort_levels.iter().any(|effort| { - open_ai::ReasoningEffort::from_str(&effort.value) - .is_ok_and(|effort| effort == open_ai::ReasoningEffort::None) - }); - - let mut request = match into_open_ai_response( - request, - &self.model.id.0, - self.model.supports_parallel_tool_calls, - true, - None, - None, - supports_none_reasoning_effort, - &OPEN_AI_PROVIDER_ID, - ) { + let request = match self.open_ai_request(request) { Ok(request) => request, Err(error) => return async move { Err(error.into()) }.boxed(), }; - if enable_thinking && let Some(effort) = effort { - request.reasoning = Some(open_ai::responses::ReasoningConfig { - effort, - summary: Some(open_ai::responses::ReasoningSummaryMode::Auto), - }); - } - let auth_context = token_provider.auth_context(cx); let executor = cx.background_executor().clone(); let future = self.request_limiter.stream(async move { @@ -1825,6 +1908,157 @@ mod tests { ); } + #[gpui::test] + async fn hosted_input_counts_use_authenticated_count_endpoint(cx: &mut gpui::TestAppContext) { + for anthropic in [true, false] { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let http_client = FakeHttpClient::create({ + let calls = calls.clone(); + move |mut request| { + let calls = calls.clone(); + async move { + calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + assert_eq!(request.uri().path(), "/count_tokens"); + assert_eq!(request.headers()["Authorization"], "Bearer test-token"); + let mut body = String::new(); + request.body_mut().read_to_string(&mut body).await?; + let body: serde_json::Value = serde_json::from_str(&body)?; + let model = if anthropic { + "claude-opus-4-6" + } else { + "gpt-5.4" + }; + assert_eq!(body["model"], model); + assert_eq!(body["provider_request"]["model"], model); + assert!(body["provider_request"].to_string().contains("aW1hZ2U=")); + Ok(Response::builder() + .status(200) + .body(AsyncBody::from(r#"{"tokens":731}"#))?) + } + } + }); + let model = if anthropic { + cloud_anthropic_test_model(http_client) + } else { + cloud_test_model(http_client) + }; + let mut request = compact_test_request(); + request.messages[0].content.push(MessageContent::Image( + language_model::LanguageModelImage { + source: "aW1hZ2U=".into(), + }, + )); + assert_eq!( + model + .count_input_tokens(request, &cx.to_async()) + .await + .unwrap(), + Some(731) + ); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + } + } + + #[gpui::test] + async fn hosted_input_counts_surface_errors_and_skip_unsupported_providers( + cx: &mut gpui::TestAppContext, + ) { + for (status, body) in [(429, "rate limit"), (200, r#"{"tokens":-1}"#), (200, "{}")] { + let model = cloud_test_model(FakeHttpClient::create(move |_| async move { + Ok(Response::builder() + .status(status) + .body(AsyncBody::from(body))?) + })); + let error = model + .count_input_tokens(compact_test_request(), &cx.to_async()) + .await + .unwrap_err(); + if status == 200 { + assert!(matches!( + error, + LanguageModelCompletionError::DeserializeResponse { .. } + )); + } else { + assert!(matches!( + error, + LanguageModelCompletionError::ProviderRejection { + category: ProviderErrorCategory::RateLimit, + .. + } + )); + } + } + let mut model = cloud_test_model(FakeHttpClient::create(|_| async { + panic!("unsupported provider must not send a counting request"); + })); + Arc::make_mut(&mut model.model).provider = cloud_llm_client::LanguageModelProvider::XAi; + assert_eq!( + model + .count_input_tokens(compact_test_request(), &cx.to_async()) + .await + .unwrap(), + None + ); + } + + #[gpui::test] + async fn retention_consent_gates_counting_compaction_and_generation( + cx: &mut gpui::TestAppContext, + ) { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let mut model = cloud_anthropic_test_model(FakeHttpClient::create({ + let calls = calls.clone(); + move |request| { + calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + async move { + assert_eq!(request.uri().path(), "/count_tokens"); + Ok(Response::builder() + .status(200) + .body(AsyncBody::from(r#"{"tokens":731}"#))?) + } + } + })); + model.id = LanguageModelId::from("claude-fable-5-1".to_string()); + Arc::make_mut(&mut model.model).id = + cloud_llm_client::LanguageModelId("claude-fable-5-1".into()); + assert!(model.requires_data_retention()); + let mut request = compact_test_request(); + request.messages[0].content.push(MessageContent::Image( + language_model::LanguageModelImage { + source: "aW1hZ2U=".into(), + }, + )); + assert!(matches!( + model + .count_input_tokens(request.clone(), &cx.to_async()) + .await, + Err(LanguageModelCompletionError::DataRetentionConsentRequired { .. }) + )); + assert!(matches!( + model.compact(request.clone(), &cx.to_async()).await, + Err(LanguageModelCompletionError::DataRetentionConsentRequired { .. }) + )); + assert!(matches!( + model + .stream_completion(request.clone(), &cx.to_async()) + .await, + Err(LanguageModelCompletionError::DataRetentionConsentRequired { .. }) + )); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 0); + + Arc::get_mut(&mut model.token_provider) + .unwrap() + .data_retention_consent = true; + assert_eq!( + model + .count_input_tokens(request, &cx.to_async()) + .await + .unwrap(), + Some(731) + ); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + } + #[gpui::test] async fn hosted_open_ai_preserves_unset_output_and_clamps_explicit_caps( cx: &mut gpui::TestAppContext, @@ -1913,7 +2147,7 @@ mod tests { is_disabled: false, disabled_reason: None, }), - token_provider: Arc::new(TestTokenProvider), + token_provider: Arc::new(TestTokenProvider::default()), http_client, app_version: None, request_limiter: RateLimiter::new(4), @@ -1945,14 +2179,17 @@ mod tests { is_disabled: false, disabled_reason: None, }), - token_provider: Arc::new(TestTokenProvider), + token_provider: Arc::new(TestTokenProvider::default()), http_client, app_version: None, request_limiter: RateLimiter::new(4), } } - struct TestTokenProvider; + #[derive(Default)] + struct TestTokenProvider { + data_retention_consent: bool, + } impl CloudLlmTokenProvider for TestTokenProvider { type AuthContext = (); @@ -1974,7 +2211,7 @@ mod tests { } fn has_data_retention_consent(&self, _cx: &impl AppContext) -> bool { - false + self.data_retention_consent } } } diff --git a/crates/open_ai/src/responses.rs b/crates/open_ai/src/responses.rs index 9c59c3bfa9..57506845dd 100644 --- a/crates/open_ai/src/responses.rs +++ b/crates/open_ai/src/responses.rs @@ -50,6 +50,18 @@ pub struct Request { } impl Request { + pub fn into_count_tokens_request(self) -> CountTokensRequest { + CountTokensRequest { + model: self.model, + instructions: self.instructions, + input: self.input, + tools: self.tools, + tool_choice: self.tool_choice, + parallel_tool_calls: self.parallel_tool_calls, + reasoning: self.reasoning, + } + } + pub fn into_compact_request(self) -> CompactRequest { CompactRequest { model: self.model, @@ -61,6 +73,23 @@ impl Request { } } +/// Structured input for counting, including images and tool definitions. +#[derive(Serialize, Debug)] +pub struct CountTokensRequest { + pub model: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub instructions: Option, + pub input: ResponseInput, + #[serde(skip_serializing_if = "Vec::is_empty")] + pub tools: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_choice: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub parallel_tool_calls: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning: Option, +} + #[derive(Serialize, Debug)] pub struct CompactRequest { pub model: String, @@ -737,9 +766,56 @@ pub async fn compact_response( request: CompactRequest, extra_headers: &CustomHeaders, ) -> Result { + send_response_request( + client, + provider_name, + api_url, + api_key, + "/responses/compact", + request, + extra_headers, + ) + .await +} + +/// Counts structured input, including images and tools, without generating output. +pub async fn count_input_tokens( + client: &dyn HttpClient, + provider_name: &str, + api_url: &str, + api_key: &str, + request: CountTokensRequest, + extra_headers: &CustomHeaders, +) -> Result { + #[derive(Deserialize)] + struct CountTokensResponse { + input_tokens: u64, + } + let response: CountTokensResponse = send_response_request( + client, + provider_name, + api_url, + api_key, + "/responses/input_tokens", + request, + extra_headers, + ) + .await?; + Ok(response.input_tokens) +} + +async fn send_response_request( + client: &dyn HttpClient, + provider_name: &str, + api_url: &str, + api_key: &str, + route: &str, + request: impl Serialize, + extra_headers: &CustomHeaders, +) -> Result { let request = HttpRequest::builder() .method(Method::POST) - .uri(format!("{api_url}/responses/compact")) + .uri(format!("{api_url}{route}")) .header("Content-Type", "application/json") .header("Authorization", format!("Bearer {}", api_key.trim())) .extra_headers(extra_headers) @@ -988,6 +1064,111 @@ mod tests { use serde_json::json; use std::sync::{Arc, Mutex}; + #[test] + fn count_input_tokens_preserves_structured_input_without_generation_fields() { + let input = json!([{ + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Describe this"}, + {"type": "input_image", "image_url": "data:image/png;base64,aW1hZ2U=", "detail": "auto"} + ] + }]); + let mut request = response_test_request(); + request.input = + ResponseInput::new(input.as_array().expect("input array").clone(), Vec::new()); + request.instructions = Some("Be precise".into()); + request.max_output_tokens = Some(1234); + request.service_tier = Some(ServiceTier::Priority); + request.prompt_cache_key = Some("thread-id".into()); + request.store = Some(false); + let generation = serde_json::to_value(&request).expect("generation request"); + assert_eq!(generation["stream"], true); + assert_eq!(generation["max_output_tokens"], 1234); + assert_eq!(generation["service_tier"], "priority"); + + let client = FakeHttpClient::create(move |mut request| { + let input = input.clone(); + async move { + assert_eq!(request.method(), Method::POST); + assert_eq!( + request.uri(), + "https://api.openai.com/v1/responses/input_tokens" + ); + assert!( + request + .headers() + .get("Authorization") + .is_some_and(|value| value == "Bearer test-key") + ); + assert_eq!(request.headers()["Content-Type"], "application/json"); + let mut body = String::new(); + request.body_mut().read_to_string(&mut body).await?; + assert_eq!( + serde_json::from_str::(&body)?, + json!({ + "model": "gpt-5.4", + "instructions": "Be precise", + "input": input + }) + ); + Ok(http_client::Response::builder() + .status(200) + .body(AsyncBody::from( + r#"{"object":"response.input_tokens","input_tokens":321}"#, + ))?) + } + }); + let count = block_on(count_input_tokens( + client.as_ref(), + "OpenAI", + "https://api.openai.com/v1", + " test-key ", + request.into_count_tokens_request(), + &CustomHeaders::default(), + )) + .expect("count succeeds"); + assert_eq!(count, 321); + } + + #[test] + fn count_input_tokens_preserves_transport_errors() { + let client = FakeHttpClient::create(|_| async { + Ok(http_client::Response::builder().status(401).body(AsyncBody::from( + r#"{"error":{"message":"Invalid API key","type":"invalid_request_error","code":"invalid_api_key"}}"# + ))?) + }); + let error = block_on(count_input_tokens( + client.as_ref(), + "OpenAI", + "https://api.openai.com/v1", + "test-key", + response_test_request().into_count_tokens_request(), + &CustomHeaders::default(), + )) + .expect_err("authentication fails"); + assert!( + matches!(error, RequestError::HttpResponseError { status_code, .. } if status_code.as_u16() == 401) + ); + + let client = FakeHttpClient::create(|_| async { + Ok(http_client::Response::builder() + .status(200) + .body(AsyncBody::from("{}"))?) + }); + assert!( + block_on(count_input_tokens( + client.as_ref(), + "OpenAI", + "https://api.openai.com/v1", + "test-key", + response_test_request().into_count_tokens_request(), + &CustomHeaders::default(), + )) + .is_err() + ); + } + #[test] fn compact_response_posts_supported_request_fields() { let captured_request = Arc::new(Mutex::new(None));