diff --git a/backend/crates/providers/openai/src/transport/request.rs b/backend/crates/providers/openai/src/transport/request.rs index 99a82daef..ecc0b23db 100644 --- a/backend/crates/providers/openai/src/transport/request.rs +++ b/backend/crates/providers/openai/src/transport/request.rs @@ -10,6 +10,8 @@ use sha2::{Digest, Sha256}; use crate::transport::protocol::responses::CodexResponsesRequest; const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers"; +const THREAD_SPAWN_SUBAGENT_KIND: &str = "thread_spawn"; +const THREAD_SPAWN_CONVERSATION_PREFIX: &str = "thread-spawn:"; const CROSS_ACCOUNT_IDENTITY_KEYS: &[&str] = &[ "authorization", @@ -160,28 +162,52 @@ fn string_value(value: Option<&Value>) -> Option { pub(crate) fn derive_conversation_anchor( request: &CodexResponsesRequest, ) -> Option<(&'static str, String)> { + thread_spawn_conversation_anchor(request).or_else(|| { + request + .client_session_id + .as_deref() + .map(|value| ("session", value.to_owned())) + .or_else(|| { + request + .client_conversation_id + .as_deref() + .map(|value| ("conversation", value.to_owned())) + }) + .or_else(|| { + request + .client_thread_id + .as_deref() + .map(|value| ("thread", value.to_owned())) + }) + .or_else(|| { + request + .prompt_cache_key() + .map(|value| ("prompt-cache", value.to_owned())) + }) + .or_else(|| derive_stable_conversation_key(request).map(|value| ("request", value))) + }) +} + +/// A `thread_spawn` shares its parent session but is a separate child execution. +/// Its child thread must therefore own the affinity and WebSocket pool identity. +fn thread_spawn_conversation_anchor( + request: &CodexResponsesRequest, +) -> Option<(&'static str, String)> { + if request.subagent_kind().as_deref() != Some(THREAD_SPAWN_SUBAGENT_KIND) { + return None; + } + request - .client_session_id + .client_thread_id .as_deref() - .map(|value| ("session", value.to_owned())) - .or_else(|| { - request - .client_conversation_id - .as_deref() - .map(|value| ("conversation", value.to_owned())) + .map(str::trim) + .filter(|thread_id| !thread_id.is_empty()) + .map(|thread_id| { + ( + "subagent-thread", + format!("{THREAD_SPAWN_CONVERSATION_PREFIX}{thread_id}"), + ) }) - .or_else(|| { - request - .client_thread_id - .as_deref() - .map(|value| ("thread", value.to_owned())) - }) - .or_else(|| { - request - .prompt_cache_key() - .map(|value| ("prompt-cache", value.to_owned())) - }) - .or_else(|| derive_stable_conversation_key(request).map(|value| ("request", value))) } const LEADING_SYSTEM_REMINDER_OPEN: &str = ""; diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index 40e96bd00..586020f41 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -212,6 +212,27 @@ fn generate_operation() -> Operation { )) } +fn generate_with_session_context( + session_id: &str, + thread_id: Option<&str>, + turn_metadata: Option<&str>, +) -> GenerateRequest { + let mut body = Map::from_iter([ + ("model".to_owned(), json!("gpt-5.4")), + ("input".to_owned(), json!("hello")), + ("session_id".to_owned(), json!(session_id)), + ]); + if let Some(thread_id) = thread_id { + body.insert("thread_id".to_owned(), json!(thread_id)); + } + if let Some(turn_metadata) = turn_metadata { + body.insert("turnMetadata".to_owned(), json!(turn_metadata)); + } + GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object("openai", body).expect("OpenAI payload"), + ) +} + fn http_generate_operation() -> Operation { let payload = ProtocolPayload::json_object( "openai", @@ -1171,6 +1192,58 @@ async fn explicit_session_id_should_override_turn_specific_prompt_cache_keys_for assert_eq!(keys[0], keys[1]); } +#[tokio::test] +async fn thread_spawn_children_should_have_independent_session_affinity_keys() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_thread_spawn_affinity").await; + let affinity = Arc::new(MemorySessionAffinity::default()); + let provider = provider_with_affinity(&store, Arc::clone(&affinity)); + let thread_spawn = r#"{"subagent_kind":"thread_spawn"}"#; + + for (request_id, thread_id, turn_metadata) in [ + ( + "req_thread_spawn_first", + Some("child-one"), + Some(thread_spawn), + ), + ( + "req_thread_spawn_second", + Some("child-two"), + Some(thread_spawn), + ), + ( + "req_thread_spawn_repeat", + Some("child-one"), + Some(thread_spawn), + ), + ("req_thread_spawn_fallback", None, Some(thread_spawn)), + ("req_thread_spawn_parent", None, None), + ] { + let stream = provider + .execute( + planned_request( + "openai", + Operation::Generate(generate_with_session_context( + "parent-session", + thread_id, + turn_metadata, + )), + ), + context(request_id, CancellationToken::new()), + ) + .await + .expect("prepare provider stream"); + drop(stream); + } + + let keys = affinity.lookup_keys(); + assert_eq!(keys.len(), 5); + assert_ne!(keys[0], keys[1]); + assert_eq!(keys[0], keys[2]); + assert_ne!(keys[3], keys[4]); + assert_ne!(keys[0], keys[3]); +} + #[tokio::test] async fn invalid_local_conversation_id_should_fall_back_to_an_opaque_affinity_key() { let store = Arc::new(MemoryAccountStore::default());