diff --git a/crates/goose-providers/src/openai.rs b/crates/goose-providers/src/openai.rs index 49a304b2d9..b7c8258d44 100644 --- a/crates/goose-providers/src/openai.rs +++ b/crates/goose-providers/src/openai.rs @@ -68,6 +68,8 @@ const DEFAULT_TIMEOUT_SECONDS: u64 = 600; type OpenAiBaseUrlParts = (String, Vec<(String, String)>, bool); +pub type OpenAiSessionIdProvider = Arc Option + Send + Sync>; + /// Ensure a base URL has an explicit scheme. /// /// Users frequently enter hosts like `localhost:1234` without a scheme. The @@ -123,7 +125,7 @@ pub fn parse_openai_base_url(raw_url: &str) -> Result { Ok((format!("{}{}", authority, path), query_params, false)) } -#[derive(Debug, serde::Serialize)] +#[derive(serde::Serialize)] pub struct OpenAiProvider { #[serde(skip)] api_client: ApiClient, @@ -139,9 +141,20 @@ pub struct OpenAiProvider { preserve_thinking_context: bool, native_openai: bool, #[serde(skip)] + session_id_provider: Option, + #[serde(skip)] n_ctx_cache: Arc>>, } +impl std::fmt::Debug for OpenAiProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OpenAiProvider") + .field("name", &self.name) + .field("base_path", &self.base_path) + .finish_non_exhaustive() + } +} + /// Builder for [`OpenAiProvider`]. /// /// Exposes every field of the provider so that constructors living outside @@ -159,6 +172,7 @@ pub struct OpenAiProviderBuilder { dynamic_models: Option, skip_canonical_filtering: bool, preserve_thinking_context: bool, + session_id_provider: Option, native_openai: bool, } @@ -176,6 +190,7 @@ impl OpenAiProviderBuilder { dynamic_models: None, skip_canonical_filtering: false, preserve_thinking_context: false, + session_id_provider: None, native_openai: false, } } @@ -248,6 +263,11 @@ impl OpenAiProviderBuilder { self } + pub fn session_id_provider(mut self, provider: OpenAiSessionIdProvider) -> Self { + self.session_id_provider = Some(provider); + self + } + pub fn native_openai(mut self, native_openai: bool) -> Self { self.native_openai = native_openai; self @@ -266,6 +286,7 @@ impl OpenAiProviderBuilder { dynamic_models: self.dynamic_models, skip_canonical_filtering: self.skip_canonical_filtering, preserve_thinking_context: self.preserve_thinking_context, + session_id_provider: self.session_id_provider, native_openai: self.native_openai, n_ctx_cache: Arc::new(Mutex::new(HashMap::new())), } @@ -304,11 +325,28 @@ impl OpenAiProvider { self.stream_responses_payload(model_config, payload).await } + fn insert_prompt_cache_key(&self, payload: &mut serde_json::Value) { + let Some(session_id) = self + .session_id_provider + .as_ref() + .and_then(|provider| provider()) + .filter(|id| !id.is_empty()) + else { + return; + }; + payload + .as_object_mut() + .unwrap() + .entry("prompt_cache_key") + .or_insert_with(|| serde_json::Value::String(session_id)); + } + async fn stream_responses_payload( &self, model_config: &ModelConfig, - payload: serde_json::Value, + mut payload: serde_json::Value, ) -> Result { + self.insert_prompt_cache_key(&mut payload); let mut log = start_log(model_config, &payload)?; let response = self .with_retry(|| async { @@ -380,6 +418,7 @@ impl OpenAiProvider { dynamic_models: None, skip_canonical_filtering: false, preserve_thinking_context: false, + session_id_provider: None, native_openai: false, n_ctx_cache: Arc::new(Mutex::new(HashMap::new())), } @@ -1053,6 +1092,7 @@ mod tests { dynamic_models: None, skip_canonical_filtering: false, preserve_thinking_context: false, + session_id_provider: None, native_openai: false, n_ctx_cache: Arc::new(Mutex::new(HashMap::new())), } @@ -1605,6 +1645,7 @@ mod tests { dynamic_models: Some(true), skip_canonical_filtering: false, preserve_thinking_context: false, + session_id_provider: None, native_openai: false, n_ctx_cache: Arc::new(Mutex::new(HashMap::new())), } diff --git a/crates/goose/src/providers/openai_def.rs b/crates/goose/src/providers/openai_def.rs index ff057be97a..de93746891 100644 --- a/crates/goose/src/providers/openai_def.rs +++ b/crates/goose/src/providers/openai_def.rs @@ -2,6 +2,7 @@ use anyhow::Result; use futures::future::BoxFuture; use goose_providers::base::ProviderDescriptor; use std::collections::HashMap; +use std::sync::Arc; use crate::config::declarative_providers::DeclarativeProviderConfig; use crate::config::Config; @@ -14,7 +15,7 @@ use crate::session_context::{ use goose_providers::api_client::{ApiClient, AuthMethod}; use goose_providers::openai::{ parse_custom_headers, parse_openai_base_url, OpenAiProvider, OpenAiProviderBuilder, - OPEN_AI_DEFAULT_BASE_PATH, OPEN_AI_VERSIONLESS_BASE_PATH, + OpenAiSessionIdProvider, OPEN_AI_DEFAULT_BASE_PATH, OPEN_AI_VERSIONLESS_BASE_PATH, }; pub struct OpenAiProviderDef; @@ -148,6 +149,7 @@ pub async fn from_env( .project(project) .custom_headers(custom_headers) .preserve_thinking_context(!is_openai) + .session_id_provider(session_id_provider()) .native_openai(is_openai) .build(); @@ -228,10 +230,15 @@ pub fn from_custom_config( None => api_client, } }) + .session_id_provider(session_id_provider()) .build() }) } +fn session_id_provider() -> OpenAiSessionIdProvider { + Arc::new(crate::session_context::current_session_id) +} + /// Components extracted from an `OPENAI_BASE_URL` value. struct ParsedBaseUrl { /// The host (scheme + authority + any path prefix before `/v1`). diff --git a/crates/goose/tests/session_id_propagation_test.rs b/crates/goose/tests/session_id_propagation_test.rs index 0666fe4489..b7a122990e 100644 --- a/crates/goose/tests/session_id_propagation_test.rs +++ b/crates/goose/tests/session_id_propagation_test.rs @@ -1,7 +1,7 @@ use goose::conversation::message::Message; use goose::providers::api_client::{ApiClient, AuthMethod}; use goose::providers::base::Provider; -use goose::providers::openai::OpenAiProvider; +use goose::providers::openai::OpenAiProviderBuilder; use goose::session_context::{session_id_request_builder, SESSION_ID_HEADER}; use goose_providers::model::ModelConfig; use serde_json::json; @@ -43,7 +43,11 @@ fn create_test_provider(mock_server_url: &str) -> Box { ) .unwrap() .with_request_builder(session_id_request_builder()); - Box::new(OpenAiProvider::new(api_client)) + Box::new( + OpenAiProviderBuilder::new(api_client) + .session_id_provider(Arc::new(goose::session_context::current_session_id)) + .build(), + ) } async fn setup_mock_server() -> (MockServer, HeaderCapture, Box) { @@ -143,8 +147,15 @@ async fn setup_mock_server() -> (MockServer, HeaderCapture, Box) { } async fn make_request(provider: &dyn Provider, session_id: &str) { + make_request_with_model(provider, session_id, ModelConfig::new("gpt-5-nano")).await; +} + +async fn make_request_with_model( + provider: &dyn Provider, + session_id: &str, + model_config: ModelConfig, +) { let message = Message::user().with_text("test message"); - let model_config = ModelConfig::new("gpt-5-nano"); let _ = goose::session_context::with_session_id( Some(session_id.to_string()), provider.complete( @@ -210,6 +221,35 @@ async fn test_session_id_propagation_to_llm() { ); } +#[tokio::test] +async fn test_session_id_is_sent_as_responses_prompt_cache_key() { + let (server, _, provider) = setup_mock_server().await; + + make_request(provider.as_ref(), "integration-test-session-123").await; + + let requests = server.received_requests().await.unwrap(); + let body: serde_json::Value = requests[0].body_json().unwrap(); + assert_eq!( + body["prompt_cache_key"], + json!("integration-test-session-123") + ); +} + +#[tokio::test] +async fn test_explicit_prompt_cache_key_overrides_session_id() { + let (server, _, provider) = setup_mock_server().await; + let model_config = + ModelConfig::new("gpt-5-nano").with_merged_request_params(std::collections::HashMap::from( + [("prompt_cache_key".to_string(), json!("shared-workflow"))], + )); + + make_request_with_model(provider.as_ref(), "session-123", model_config).await; + + let requests = server.received_requests().await.unwrap(); + let body: serde_json::Value = requests[0].body_json().unwrap(); + assert_eq!(body["prompt_cache_key"], json!("shared-workflow")); +} + #[tokio::test] async fn test_session_id_always_present() { let (_, capture, provider) = setup_mock_server().await;