From cb7f69850bddae81a35add62003987028f80b6b5 Mon Sep 17 00:00:00 2001 From: Zyy Date: Fri, 28 Aug 2026 02:20:55 +0800 Subject: [PATCH] feat(openai): add session affinity observability --- .../openai/src/credential/affinity.rs | 76 ++++++++- .../providers/openai/src/credential/mod.rs | 2 +- .../openai/src/credential/selector.rs | 18 ++ .../crates/providers/openai/src/provider.rs | 14 +- .../openai/tests/provider/contract.rs | 159 +++++++++++++++++- 5 files changed, 254 insertions(+), 15 deletions(-) diff --git a/backend/crates/providers/openai/src/credential/affinity.rs b/backend/crates/providers/openai/src/credential/affinity.rs index 3c03727c2..706aa7e32 100644 --- a/backend/crates/providers/openai/src/credential/affinity.rs +++ b/backend/crates/providers/openai/src/credential/affinity.rs @@ -1,4 +1,4 @@ -//! OpenAI 根会话锚点到 Store 不透明账号亲和键的单向派生。 +//! OpenAI 根会话锚点到 Store 不透明账号亲和键及诊断上下文的单向派生。 use gateway_core::policy::ClientApiKeyId; use gateway_core::provider_ports::ProviderSessionAffinityKey; @@ -7,20 +7,82 @@ use sha2::{Digest, Sha256}; use crate::transport::protocol::responses::CodexResponsesRequest; use crate::transport::request::derive_conversation_anchor; -pub(crate) fn derive_codex_session_affinity_key( +const AFFINITY_KEY_HASH_LENGTH: usize = 12; + +/// 一次请求派生出的账号亲和键及其结构化日志上下文。 +pub(crate) struct CodexSessionAffinity { + key: ProviderSessionAffinityKey, + key_hash: String, + anchor_source: &'static str, + anchor: String, + session_id: Option, +} + +impl CodexSessionAffinity { + #[must_use] + pub(crate) const fn key(&self) -> &ProviderSessionAffinityKey { + &self.key + } + + #[must_use] + pub(crate) fn key_hash(&self) -> &str { + &self.key_hash + } + + #[must_use] + pub(crate) const fn anchor_source(&self) -> &'static str { + self.anchor_source + } + + #[must_use] + pub(crate) fn anchor(&self) -> &str { + &self.anchor + } + + #[must_use] + pub(crate) fn session_id(&self) -> Option<&str> { + self.session_id.as_deref() + } + + #[must_use] + pub(crate) const fn session_id_present(&self) -> bool { + self.session_id.is_some() + } + + #[must_use] + pub(crate) fn into_key(self) -> ProviderSessionAffinityKey { + self.key + } +} + +pub(crate) fn derive_codex_session_affinity( request: &CodexResponsesRequest, client_api_key_id: &ClientApiKeyId, -) -> Option { - let (domain, value) = derive_account_affinity_anchor(request)?; - let session_key = opaque_affinity_key(domain, &value)?; - opaque_affinity_key( +) -> Option { + let session_id = non_empty(request.client_session_id.as_deref()).map(str::to_owned); + let (anchor_source, anchor) = derive_account_affinity_anchor(request)?; + let session_key = opaque_affinity_key(anchor_source, &anchor)?; + let key = opaque_affinity_key( "client-session", &format!( "{}\0{}", client_api_key_id.as_str(), session_key.expose_to_store() ), - ) + )?; + // 亲和键本身已经是 SHA-256；日志沿用 WebSocket 诊断的 12 位短哈希长度。 + let key_hash = key + .expose_to_store() + .chars() + .take(AFFINITY_KEY_HASH_LENGTH) + .collect(); + Some(CodexSessionAffinity { + key, + key_hash, + anchor_source, + anchor, + session_id, + }) } /// 账号选择优先使用根会话事实；`local_conversation_id` 仍可按 child thread 隔离 diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index fe831c05a..59da4d6a9 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -15,7 +15,7 @@ pub mod token_client; mod types; pub(crate) use affinity::{ - derive_codex_cyber_policy_session_key, derive_codex_session_affinity_key, + CodexSessionAffinity, derive_codex_cyber_policy_session_key, derive_codex_session_affinity, }; pub(crate) use oauth::oauth_owner_ref; pub(crate) use types::parse_access_token_expiration; diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index cc16c98dd..1599b0a16 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -22,6 +22,7 @@ use secrecy::ExposeSecret; use thiserror::Error; use url::Url; +use super::affinity::CodexSessionAffinity; use super::catalog::CodexCredentialCatalogService; use super::cookie::CodexCookiePolicy; use super::quota::CodexCredentialQuotaService; @@ -102,6 +103,7 @@ struct CredentialSelectionInput<'a> { request_url: &'a Url, attempt: &'a AttemptContext, session_affinity_key: Option<&'a ProviderSessionAffinityKey>, + session_affinity_observation: Option<&'a CodexSessionAffinity>, } #[derive(Clone)] @@ -291,6 +293,7 @@ impl CodexCredentialSelector { request_url: request.request_url, attempt: request.attempt, session_affinity_key: request.session_affinity_key, + session_affinity_observation: None, }; self.select_inner( &input, @@ -304,11 +307,13 @@ impl CodexCredentialSelector { &self, request: &SelectCodexCredential<'_>, cyber_policy_session_key: Option<&ProviderSessionAffinityKey>, + session_affinity_observation: Option<&CodexSessionAffinity>, ) -> Result { let input = CredentialSelectionInput { request_url: request.request_url, attempt: request.attempt, session_affinity_key: request.session_affinity_key, + session_affinity_observation, }; self.select_inner( &input, @@ -330,6 +335,7 @@ impl CodexCredentialSelector { request_url: request.request_url, attempt: request.attempt, session_affinity_key: None, + session_affinity_observation: None, }; self.select_inner(&input, None, ModelCatalogEligibility::NotApplicable) .await @@ -563,6 +569,7 @@ impl CodexCredentialSelector { .unwrap_or_else(|| account.id().clone()) }; let affinity_telemetry = affinity.telemetry(account.id()); + let affinity_observation = request.session_affinity_observation; tracing::info!( request_id = %request.attempt.request_id(), attempt_index = request.attempt.attempt_index().get(), @@ -573,6 +580,17 @@ impl CodexCredentialSelector { .escape_reason .map_or("", AffinityEscapeReason::as_str), account_switch = affinity_telemetry.account_switch, + affinity_key_hash = affinity_observation + .map_or("", CodexSessionAffinity::key_hash), + affinity_anchor_source = affinity_observation + .map_or("", CodexSessionAffinity::anchor_source), + affinity_anchor = affinity_observation + .map_or("", CodexSessionAffinity::anchor), + session_id = affinity_observation + .and_then(CodexSessionAffinity::session_id) + .unwrap_or(""), + session_id_present = affinity_observation + .is_some_and(CodexSessionAffinity::session_id_present), "OpenAI account selected" ); let runtime = self.repository.load_runtime_credential(&account).await?; diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs index a629faf06..604a57d87 100644 --- a/backend/crates/providers/openai/src/provider.rs +++ b/backend/crates/providers/openai/src/provider.rs @@ -52,9 +52,9 @@ use crate::credential::{ CodexAccountFailure, CodexCredentialCatalogError, CodexCredentialCatalogService, CodexCredentialLease, CodexCredentialQuotaService, CodexCredentialRefreshOutcome, CodexCredentialRefreshService, CodexCredentialSelector, CodexCyberPolicyScope, - CodexQuotaRefreshPolicy, CredentialSelectionError, RuntimeCodexCookie, SelectCodexCredential, - SelectCodexProviderEndpointCredential, derive_codex_cyber_policy_session_key, - derive_codex_session_affinity_key, + CodexQuotaRefreshPolicy, CodexSessionAffinity, CredentialSelectionError, RuntimeCodexCookie, + SelectCodexCredential, SelectCodexProviderEndpointCredential, + derive_codex_cyber_policy_session_key, derive_codex_session_affinity, }; use crate::transport::canonical::{ CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, @@ -285,8 +285,8 @@ impl Provider for CodexProvider { .as_ref() .and_then(|state| state.turn_state.clone()); } - let session_affinity_key = - derive_codex_session_affinity_key(&upstream_request, context.client_api_key_ref()); + let session_affinity = + derive_codex_session_affinity(&upstream_request, context.client_api_key_ref()); let cyber_policy_session_key = derive_codex_cyber_policy_session_key(&upstream_request, context.client_api_key_ref()); let transport = selected_transport(&upstream_request); @@ -299,9 +299,10 @@ impl Provider for CodexProvider { upstream_model: upstream_model.as_str(), request_url: &self.responses_url, attempt: &context, - session_affinity_key: session_affinity_key.as_ref(), + session_affinity_key: session_affinity.as_ref().map(|affinity| affinity.key()), }, cyber_policy_session_key.as_ref(), + session_affinity.as_ref(), ) .await .map_err(map_selection_error)?; @@ -404,6 +405,7 @@ impl Provider for CodexProvider { continuation_scope: None, }); let allows_account_state_mutation = lease.allows_account_state_mutation(); + let session_affinity_key = session_affinity.map(CodexSessionAffinity::into_key); let events = cold_response_stream(ColdResponse { client: self.client.clone(), response_origin: self.responses_url.clone(), diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index 290cb9710..89d14e825 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -1,6 +1,7 @@ use std::collections::{BTreeMap, BTreeSet}; +use std::io::{self, Write}; use std::num::NonZeroU32; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; use base64::{Engine as _, engine::general_purpose::STANDARD}; @@ -48,6 +49,7 @@ use tokio::{ }; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::protocol::{CloseFrame, frame::coding::CloseCode}; +use tracing_subscriber::fmt::MakeWriter; use wiremock::matchers::{body_bytes, header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -64,6 +66,61 @@ const CAPTURE_COMPLETED_SSE: &str = concat!( "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_scope_capture\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n" ); +#[derive(Clone, Default)] +struct CapturedLogs { + bytes: Arc>>, +} + +impl CapturedLogs { + fn json_events(&self) -> Vec { + let bytes = self.bytes.lock().expect("captured logs lock").clone(); + String::from_utf8(bytes) + .expect("captured logs are UTF-8") + .lines() + .map(|line| serde_json::from_str(line).expect("captured log is JSON")) + .collect() + } +} + +impl<'writer> MakeWriter<'writer> for CapturedLogs { + type Writer = Self; + + fn make_writer(&'writer self) -> Self::Writer { + self.clone() + } +} + +impl Write for CapturedLogs { + fn write(&mut self, buffer: &[u8]) -> io::Result { + self.bytes + .lock() + .map_err(|_| io::Error::other("captured logs lock poisoned"))? + .extend_from_slice(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +fn selected_account_log_fields<'events>( + events: &'events [Value], + request_id: &str, +) -> &'events Map { + events + .iter() + .find_map(|event| { + let fields = event.get("fields")?.as_object()?; + (fields.get("message").and_then(Value::as_str) == Some("OpenAI account selected") + && fields.get("request_id").and_then(Value::as_str) == Some(request_id)) + .then_some(fields) + }) + .unwrap_or_else(|| { + panic!("OpenAI account selection log for {request_id}; captured events: {events:#?}") + }) +} + fn wire_profile() -> CodexWireProfileState { CodexWireProfileState::new(CodexWireProfile { originator: "codex_cli_rs".to_owned(), @@ -1740,6 +1797,106 @@ async fn account_change_drops_only_account_bound_opaque_headers() { ); } +#[tokio::test(flavor = "current_thread")] +async fn account_selection_log_should_include_affinity_observation_fields() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_affinity").await; + let provider = provider_with_affinity(&store, Arc::new(MemorySessionAffinity::default())); + let captured = CapturedLogs::default(); + let subscriber = tracing_subscriber::fmt() + .json() + .without_time() + .with_ansi(false) + .with_target(false) + .with_writer(captured.clone()) + .finish(); + // 该 integration test binary 没有其他 subscriber；全局安装可避免并行测试切换 + // thread-local dispatcher 时重建 tracing callsite interest 所产生的竞争。 + tracing::subscriber::set_global_default(subscriber) + .expect("install affinity observation log subscriber"); + + for (request_id, prompt_cache_key, session_id) in [ + ( + "req_affinity_observation_session_first", + "turn-cache-first", + Some("stable-observation-session"), + ), + ( + "req_affinity_observation_session_second", + "turn-cache-second", + Some("stable-observation-session"), + ), + ( + "req_affinity_observation_conversation", + "turn-cache-with-conversation", + None, + ), + ] { + let mut payload = ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("gpt-5.4")), + ("input".to_owned(), json!("hello")), + ("prompt_cache_key".to_owned(), json!(prompt_cache_key)), + ]), + ) + .expect("OpenAI payload"); + if let Some(session_id) = session_id { + payload = payload.with_context(Map::from_iter([( + "session_id".to_owned(), + json!(session_id), + )])); + } else { + payload = payload.with_context(Map::from_iter([( + "conversation_id".to_owned(), + json!("root-observation-conversation"), + )])); + } + let generation = GenerateRequest::from_protocol_payload(payload); + let stream = provider + .execute( + planned_request("openai", Operation::Generate(generation)), + context(request_id, CancellationToken::new()), + ) + .await + .expect("prepare affinity observation request"); + drop(stream); + } + + let events = captured.json_events(); + let first = selected_account_log_fields(&events, "req_affinity_observation_session_first"); + let second = selected_account_log_fields(&events, "req_affinity_observation_session_second"); + for fields in [first, second] { + assert_eq!(fields["affinity_anchor_source"], "root-session"); + assert_eq!(fields["affinity_anchor"], "stable-observation-session"); + assert_eq!(fields["session_id"], "stable-observation-session"); + assert_eq!(fields["session_id_present"], true); + let key_hash = fields["affinity_key_hash"] + .as_str() + .expect("affinity key hash"); + assert_eq!(key_hash.len(), 12); + assert!(key_hash.bytes().all(|byte| byte.is_ascii_hexdigit())); + } + assert_eq!( + first["affinity_key_hash"], second["affinity_key_hash"], + "the same session must emit the same affinity hash" + ); + + let conversation = + selected_account_log_fields(&events, "req_affinity_observation_conversation"); + assert_eq!(conversation["affinity_anchor_source"], "root-conversation"); + assert_eq!( + conversation["affinity_anchor"], + "root-observation-conversation" + ); + assert_eq!(conversation["session_id"], ""); + assert_eq!(conversation["session_id_present"], false); + assert_ne!( + first["affinity_key_hash"], conversation["affinity_key_hash"], + "different anchors must not share the same affinity hash" + ); +} + #[tokio::test] async fn prompt_cache_key_should_become_an_opaque_session_affinity_lookup_key() { let store = Arc::new(MemoryAccountStore::default());