From 0dcb2c3ce1fb16fb585708b2c20f2f91d765fbb3 Mon Sep 17 00:00:00 2001 From: Bennet Bo Fenner Date: Fri, 7 Aug 2026 12:37:14 +0200 Subject: [PATCH] open_ai: Implement server-side compaction for subscription (#62313) Release Notes: - N/A --- crates/open_ai/src/responses.rs | 42 +++- .../src/openai_subscribed.rs | 238 ++++++++++++++---- 2 files changed, 228 insertions(+), 52 deletions(-) diff --git a/crates/open_ai/src/responses.rs b/crates/open_ai/src/responses.rs index dc2d6e8dec..b166f762dc 100644 --- a/crates/open_ai/src/responses.rs +++ b/crates/open_ai/src/responses.rs @@ -75,10 +75,14 @@ pub struct CompactRequest { #[derive(Deserialize, Debug)] pub struct CompactedResponse { + #[serde(default)] pub id: String, + #[serde(default)] pub created_at: u64, + #[serde(default)] pub object: String, pub output: Vec, + #[serde(default)] pub usage: ResponseUsage, } @@ -145,10 +149,21 @@ fn validate_compaction_items(items: &[Value]) -> Result<()> { if !items.iter().any(|item| { item.get("type") .and_then(Value::as_str) - .is_some_and(|item_type| item_type == "compaction") + .is_some_and(|item_type| { + matches!( + item_type, + "compaction" | "compaction_summary" | "context_compaction" + ) + }) }) { + let item_types = items + .iter() + .filter_map(|item| item.get("type").and_then(Value::as_str)) + .collect::>() + .join(", "); return Err(anyhow!( - "OpenAI compaction output did not contain a compaction item" + "OpenAI compaction output did not contain a recognized compaction item \ + (output types: {item_types})" )); } Ok(()) @@ -1129,6 +1144,29 @@ mod tests { ); } + #[test] + fn compacted_response_preserves_legacy_compaction_item_types() { + for item_type in ["compaction_summary", "context_compaction"] { + let output = vec![json!({ + "type": item_type, + "encrypted_content": "opaque-state" + })]; + let response: CompactedResponse = + serde_json::from_value(json!({ "output": &output })).unwrap(); + + let CompactedContext::ProviderState(state) = response + .into_compacted_context(OPEN_AI_PROVIDER_ID) + .unwrap() + else { + panic!("expected provider state"); + }; + assert_eq!( + provider_compaction_items(&state, &OPEN_AI_PROVIDER_ID).unwrap(), + Some(output) + ); + } + } + #[test] fn compacted_response_rejects_output_without_compaction_item() { let response: CompactedResponse = serde_json::from_value(json!({ diff --git a/crates/openai_subscribed/src/openai_subscribed.rs b/crates/openai_subscribed/src/openai_subscribed.rs index 21473f05b0..c6464a5bf9 100644 --- a/crates/openai_subscribed/src/openai_subscribed.rs +++ b/crates/openai_subscribed/src/openai_subscribed.rs @@ -9,11 +9,14 @@ use http_client::{ http::{HeaderName, HeaderValue}, }; use language_model::{ - LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, + CompactionResult, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelEffortLevel, LanguageModelId, LanguageModelName, LanguageModelProviderId, LanguageModelProviderName, LanguageModelRequest, LanguageModelToolChoice, RateLimiter, }; -use open_ai::{ReasoningEffort, responses::stream_response}; +use open_ai::{ + ReasoningEffort, + responses::{compact_response, stream_response}, +}; use rand::RngCore as _; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; @@ -22,7 +25,9 @@ use std::time::{SystemTime, UNIX_EPOCH}; use url::form_urlencoded; use util::ResultExt as _; -use open_ai::completion::{OpenAiResponseEventMapper, into_open_ai_response}; +use open_ai::completion::{ + OpenAiResponseEventMapper, into_open_ai_response, token_usage_from_response_usage, +}; pub const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("openai-subscribed"); pub const PROVIDER_NAME: LanguageModelProviderName = @@ -374,6 +379,52 @@ struct OpenAiSubscribedLanguageModel { request_limiter: RateLimiter, } +impl OpenAiSubscribedLanguageModel { + fn codex_responses_request( + &self, + mut request: LanguageModelRequest, + ) -> Result { + if !self.model.supports_priority() { + request.speed = None; + } + let mut responses_request = into_open_ai_response( + request, + self.model.id(), + self.model.supports_parallel_tool_calls(), + self.model.supports_prompt_cache_key(), + None, + self.model.default_reasoning_effort(), + self.model + .supported_reasoning_efforts() + .contains(&ReasoningEffort::None), + &PROVIDER_ID, + )?; + responses_request.store = Some(false); + responses_request.instructions.get_or_insert_default(); + Ok(responses_request) + } +} + +fn codex_extra_headers(credentials: &CodexCredentials) -> CustomHeaders { + let mut header_pairs: Vec<(HeaderName, HeaderValue)> = vec![ + ( + HeaderName::from_static("originator"), + HeaderValue::from_static("zed"), + ), + ( + HeaderName::from_static("openai-beta"), + HeaderValue::from_static("responses=experimental"), + ), + ]; + if let Some(id) = &credentials.account_id + && !id.is_empty() + && let Ok(value) = HeaderValue::from_str(id) + { + header_pairs.push((HeaderName::from_static("chatgpt-account-id"), value)); + } + CustomHeaders::new(header_pairs) +} + impl LanguageModel for OpenAiSubscribedLanguageModel { fn id(&self) -> LanguageModelId { self.id.clone() @@ -415,6 +466,49 @@ impl LanguageModel for OpenAiSubscribedLanguageModel { self.model.supports_priority() } + fn supports_explicit_compaction(&self) -> bool { + true + } + + fn compact( + &self, + request: LanguageModelRequest, + cx: &AsyncApp, + ) -> BoxFuture<'static, Result> { + let compact_request = match self.codex_responses_request(request) { + Ok(responses_request) => responses_request.into_compact_request(), + Err(error) => return async move { Err(error.into()) }.boxed(), + }; + + let state = self.state.downgrade(); + let http_client = self.http_client.clone(); + let request_limiter = self.request_limiter.clone(); + + cx.spawn(async move |cx| { + let creds = get_fresh_credentials(&state, &http_client, cx).await?; + let extra_headers = codex_extra_headers(&creds); + request_limiter + .run(async move { + let response = compact_response( + http_client.as_ref(), + PROVIDER_NAME.0.as_str(), + CODEX_BASE_URL, + &creds.access_token, + compact_request, + &extra_headers, + ) + .await?; + let usage = token_usage_from_response_usage(&response.usage); + let context = response + .into_compacted_context(PROVIDER_ID) + .map_err(LanguageModelCompletionError::Other)?; + Ok(CompactionResult { context, usage }) + }) + .await + }) + .boxed() + } + fn supported_effort_levels(&self) -> Vec { let default_effort = self.model.default_reasoning_effort(); self.model @@ -455,7 +549,7 @@ impl LanguageModel for OpenAiSubscribedLanguageModel { fn stream_completion( &self, - mut request: LanguageModelRequest, + request: LanguageModelRequest, cx: &AsyncApp, ) -> BoxFuture< 'static, @@ -467,35 +561,10 @@ impl LanguageModel for OpenAiSubscribedLanguageModel { LanguageModelCompletionError, >, > { - if !self.model.supports_priority() { - request.speed = None; - } - - // The Codex backend rejects `max_output_tokens` (`Unsupported parameter`), - // unlike the public OpenAI Responses API. Pass `None` so the field is - // omitted from the serialized request body entirely. - let mut responses_request = match into_open_ai_response( - request, - self.model.id(), - self.model.supports_parallel_tool_calls(), - self.model.supports_prompt_cache_key(), - /*max_output_tokens*/ None, - self.model.default_reasoning_effort(), - self.model - .supported_reasoning_efforts() - .contains(&ReasoningEffort::None), - &PROVIDER_ID, - ) { - Ok(request) => request, + let responses_request = match self.codex_responses_request(request) { + Ok(responses_request) => responses_request, Err(error) => return async move { Err(error.into()) }.boxed(), }; - responses_request.store = Some(false); - - // `into_open_ai_response` already hoists system messages into - // `instructions`, which is the only form the Codex backend accepts. - // Codex has only ever been sent requests with the field present - // (possibly empty), so keep sending it even without system messages. - responses_request.instructions.get_or_insert_default(); let state = self.state.downgrade(); let http_client = self.http_client.clone(); @@ -503,25 +572,7 @@ impl LanguageModel for OpenAiSubscribedLanguageModel { let future = cx.spawn(async move |cx| { let creds = get_fresh_credentials(&state, &http_client, cx).await?; - - let mut header_pairs: Vec<(HeaderName, HeaderValue)> = vec![ - ( - HeaderName::from_static("originator"), - HeaderValue::from_static("zed"), - ), - ( - HeaderName::from_static("openai-beta"), - HeaderValue::from_static("responses=experimental"), - ), - ]; - if let Some(ref id) = creds.account_id { - if !id.is_empty() { - if let Ok(value) = HeaderValue::from_str(id) { - header_pairs.push((HeaderName::from_static("chatgpt-account-id"), value)); - } - } - } - let extra_headers = CustomHeaders::new(header_pairs); + let extra_headers = codex_extra_headers(&creds); let access_token = creds.access_token.clone(); request_limiter @@ -1256,6 +1307,93 @@ mod tests { }); } + #[gpui::test] + async fn test_compact_posts_to_codex_compact_endpoint(cx: &mut TestAppContext) { + let compact_request_count = Arc::new(AtomicUsize::new(0)); + let http_client = FakeHttpClient::create({ + let compact_request_count = compact_request_count.clone(); + move |request| { + let compact_request_count = compact_request_count.clone(); + async move { + assert_eq!( + request.uri().to_string(), + "https://chatgpt.com/backend-api/codex/responses/compact" + ); + assert_eq!( + request + .headers() + .get("authorization") + .and_then(|value| value.to_str().ok()), + Some("Bearer fresh_access") + ); + assert_eq!( + request + .headers() + .get("chatgpt-account-id") + .and_then(|value| value.to_str().ok()), + Some("account-123") + ); + compact_request_count.fetch_add(1, Ordering::SeqCst); + // Only `output` is present, matching the least metadata a + // Responses-protocol backend is known to attach. + Ok(http_client::Response::builder().status(200).body( + http_client::AsyncBody::from( + serde_json::json!({ + "output": [{ + "type": "compaction", + "id": "cmp_1", + "encrypted_content": "opaque-state", + }], + }) + .to_string(), + ), + )?) + } + } + }); + + let http: Arc = http_client; + let mut credentials = make_fresh_credentials(); + credentials.account_id = Some("account-123".to_string()); + let state = make_state(http, Some(credentials), cx); + let model = cx.read(|cx| create_language_model(ChatGptModel::Gpt55, &state, cx)); + assert!(model.supports_explicit_compaction()); + + let request = LanguageModelRequest { + messages: vec![language_model::LanguageModelRequestMessage { + role: language_model::Role::User, + content: vec![language_model::MessageContent::Text("Hello".into())], + cache: false, + reasoning_details: None, + }], + ..Default::default() + }; + let async_cx = cx.to_async(); + let result = model + .compact(request, &async_cx) + .await + .expect("compaction should succeed"); + + assert_eq!(compact_request_count.load(Ordering::SeqCst), 1); + assert_eq!(result.usage, language_model::TokenUsage::default()); + let language_model::CompactedContext::ProviderState(compaction_state) = result.context + else { + panic!("expected provider compaction state"); + }; + assert_eq!(compaction_state.provider_id(), &PROVIDER_ID); + let items = open_ai::responses::provider_compaction_items(&compaction_state, &PROVIDER_ID) + .expect("the compacted state should parse") + .expect("the compacted state should be owned by the subscription provider"); + assert_eq!( + items, + vec![serde_json::json!({ + "type": "compaction", + "id": "cmp_1", + "encrypted_content": "opaque-state", + })] + ); + } + struct FakeCredentialsProvider { storage: Mutex)>>, }