fix(gateway): 对齐 Responses WebSocket 前导恢复与续接语义

crates/service/src/agent_identity.rs

@@ -507,6 +507,12 @@ fn register_agent_identity(
     let mut request_builder = client
         .post(&url)
         .timeout(AGENT_REGISTRATION_TIMEOUT)
+        // Keep the auth request on the same first-party identity path as the
+        // official Codex client.  The account bearer token alone is valid for
+        // normal Responses calls, but agent registration also relies on the
+        // Codex originator and User-Agent headers for request classification.
+        .header("originator", crate::gateway::current_wire_originator())
+        .header("User-Agent", crate::gateway::current_codex_user_agent())
         .bearer_auth(access_token)
         .json(&request);
     if is_fedramp {
@@ -671,6 +677,10 @@ fn register_agent_identity_task(
     let response = client
         .post(&url)
         .timeout(AGENT_TASK_REGISTRATION_TIMEOUT)
+        // The official Codex auth client applies the same identity headers to
+        // task registration as to runtime registration.
+        .header("originator", crate::gateway::current_wire_originator())
+        .header("User-Agent", crate::gateway::current_codex_user_agent())
         .json(&request)
         .send()
         .map_err(|err| format!("agent task registration request failed: {err}"))?;
@@ -1011,13 +1021,15 @@ mod tests {
                     .expect("registration request");
                 let path = request.url().to_string();
                 let authorization = request_header(&request, "authorization");
+                let originator = request_header(&request, "originator");
+                let user_agent = request_header(&request, "user-agent");
                 let mut body = String::new();
                 request
                     .as_reader()
                     .read_to_string(&mut body)
                     .expect("read registration request");
                 request_tx
-                    .send((path, authorization, body))
+                    .send((path, authorization, originator, user_agent, body))
                     .expect("record request");
                 request
                     .respond(
@@ -1050,10 +1062,16 @@ mod tests {
         );
         assert!(authorization.value.starts_with("AgentAssertion "));
 
-        let (registration_path, registration_auth, registration_body) = request_rx
+        let (
+            registration_path,
+            registration_auth,
+            registration_originator,
+            registration_user_agent,
+            registration_body,
+        ) = request_rx
             .recv_timeout(Duration::from_secs(5))
             .expect("receive identity registration");
-        let (task_path, task_auth, _task_body) = request_rx
+        let (task_path, task_auth, _task_originator, _task_user_agent, _task_body) = request_rx
             .recv_timeout(Duration::from_secs(5))
             .expect("receive task registration");
         server_handle.join().expect("join registration server");
@@ -1062,6 +1080,14 @@ mod tests {
             registration_auth.as_deref(),
             Some(format!("Bearer {access_token}").as_str())
         );
+        assert_eq!(
+            registration_originator.as_deref(),
+            Some(crate::gateway::current_wire_originator().as_str())
+        );
+        assert_eq!(
+            registration_user_agent.as_deref(),
+            Some(crate::gateway::current_codex_user_agent().as_str())
+        );
         let registration_body: serde_json::Value =
             serde_json::from_str(&registration_body).expect("parse registration body");
         assert_eq!(
@@ -1363,12 +1389,16 @@ mod tests {
                 .expect("registration server timeout")
                 .expect("registration request");
             let path = request.url().to_string();
+            let originator = request_header(&request, "originator");
+            let user_agent = request_header(&request, "user-agent");
             let mut body = String::new();
             request
                 .as_reader()
                 .read_to_string(&mut body)
                 .expect("read registration request");
-            request_tx.send((path, body)).expect("record request");
+            request_tx
+                .send((path, originator, user_agent, body))
+                .expect("record request");
             request
                 .respond(
                     Response::from_string(r#"{"taskId":"task-from-server"}"#)
@@ -1385,11 +1415,19 @@ mod tests {
             register_agent_identity_task(&reqwest::blocking::Client::new(), &identity, &base_url)
                 .expect("register task");
         assert_eq!(task_id, "task-from-server");
-        let (path, body) = request_rx
+        let (path, originator, user_agent, body) = request_rx
             .recv_timeout(Duration::from_secs(5))
             .expect("receive request");
         server_handle.join().expect("join server");
         assert_eq!(path, "/v1/agent/agent-runtime-1/task/register");
+        assert_eq!(
+            originator.as_deref(),
+            Some(crate::gateway::current_wire_originator().as_str())
+        );
+        assert_eq!(
+            user_agent.as_deref(),
+            Some(crate::gateway::current_codex_user_agent().as_str())
+        );
         let body: serde_json::Value = serde_json::from_str(&body).expect("parse request body");
         let timestamp = body["timestamp"].as_str().expect("timestamp");
         let signature = BASE64_STANDARD

crates/service/src/gateway/auth/tests/token_exchange_tests.rs

@@ -1,4 +1,5 @@
 use super::*;
+use base64::Engine as _;
 
 /// 函数 `same_account_reuses_exchange_lock`
 ///
@@ -115,7 +116,7 @@ fn fallback_to_access_token_uses_runtime_access_token_when_exchange_fails() {
     assert_eq!(bearer, "runtime-access-token");
 }
 
-/// 函数 `api_key_exchange_subject_tokens_falls_back_to_imported_access_token`
+/// 函数 `api_key_exchange_subject_token_omits_access_token_without_id_token`
 ///
 /// 作者: gaohongshun
 ///
@@ -127,7 +128,7 @@ fn fallback_to_access_token_uses_runtime_access_token_when_exchange_fails() {
 /// # 返回
 /// 无
 #[test]
-fn api_key_exchange_subject_tokens_falls_back_to_imported_access_token() {
+fn api_key_exchange_subject_token_omits_access_token_without_id_token() {
     let token = Token {
         account_id: "acc-import-session".to_string(),
         id_token: String::new(),
@@ -137,13 +138,10 @@ fn api_key_exchange_subject_tokens_falls_back_to_imported_access_token() {
         last_refresh: now_ts(),
     };
 
-    assert_eq!(
-        api_key_exchange_subject_tokens(&token),
-        vec!["imported-session-access".to_string()]
-    );
+    assert_eq!(api_key_exchange_subject_token(&token), None);
 }
 
-/// 函数 `api_key_exchange_subject_tokens_prefers_access_token_before_id_token`
+/// 函数 `api_key_exchange_subject_token_uses_id_token_only`
 ///
 /// 作者: gaohongshun
 ///
@@ -155,7 +153,7 @@ fn api_key_exchange_subject_tokens_falls_back_to_imported_access_token() {
 /// # 返回
 /// 无
 #[test]
-fn api_key_exchange_subject_tokens_prefers_access_token_before_id_token() {
+fn api_key_exchange_subject_token_uses_id_token_only() {
     let token = Token {
         account_id: "acc-login".to_string(),
         id_token: "id-token".to_string(),
@@ -166,8 +164,32 @@ fn api_key_exchange_subject_tokens_prefers_access_token_before_id_token() {
     };
 
     assert_eq!(
-        api_key_exchange_subject_tokens(&token),
-        vec!["access-token".to_string(), "id-token".to_string()]
+        api_key_exchange_subject_token(&token),
+        Some("id-token".to_string())
+    );
+}
+
+#[test]
+fn api_key_exchange_client_id_prefers_id_token_claim() {
+    let jwt = |client_id: &str| {
+        let payload = serde_json::json!({"sub":"user-test","client_id": client_id}).to_string();
+        format!(
+            "header.{}.signature",
+            base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(payload)
+        )
+    };
+    let token = Token {
+        account_id: "acc-client-id".to_string(),
+        id_token: jwt("id-token-client"),
+        access_token: jwt("access-token-client"),
+        refresh_token: String::new(),
+        api_key_access_token: None,
+        last_refresh: now_ts(),
+    };
+
+    assert_eq!(
+        api_key_exchange_client_id(&token, "fallback-client"),
+        "id-token-client"
     );
 }
 

crates/service/src/gateway/auth/token_exchange.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex, OnceLock};
 
-use codexmanager_core::auth::extract_token_exp;
+use codexmanager_core::auth::{extract_client_id_claim, extract_token_exp, DEFAULT_CLIENT_ID};
 use codexmanager_core::storage::{now_ts, Account, Storage, Token};
 
 use crate::account_status::mark_account_unavailable_for_auth_error;
@@ -139,40 +139,39 @@ fn exchange_and_persist_api_key_access_token(
     issuer: &str,
     client_id: &str,
 ) -> Result<String, String> {
-    let mut errors = Vec::new();
-    for subject_token in api_key_exchange_subject_tokens(token) {
-        match auth_tokens::obtain_api_key(issuer, client_id, &subject_token) {
-            Ok(exchanged) => {
-                token.api_key_access_token = Some(exchanged.clone());
-                let _ = storage.insert_token(token);
-                return Ok(exchanged);
-            }
-            Err(err) => errors.push(err),
+    let Some(subject_token) = api_key_exchange_subject_token(token) else {
+        return Err("id_token is unavailable for API key token exchange".to_string());
+    };
+    match auth_tokens::obtain_api_key(issuer, client_id, &subject_token) {
+        Ok(exchanged) => {
+            token.api_key_access_token = Some(exchanged.clone());
+            let _ = storage.insert_token(token);
+            Ok(exchanged)
         }
+        Err(err) => Err(err),
     }
-
-    let exchange_error = errors
-        .into_iter()
-        .next()
-        .unwrap_or_else(|| "api key exchange subject token is missing".to_string());
-    Err(exchange_error)
 }
 
-fn api_key_exchange_subject_tokens(token: &Token) -> Vec<String> {
-    let mut subjects = Vec::new();
-    // 中文注释：直接导入 api/auth/session JSON 时通常只有 accessToken；
-    // 登录授权路径已优先缓存 api_key_access_token，缓存缺失时也先按最新 access_token 兑换。
-    push_unique_subject_token(&mut subjects, token.access_token.as_str());
-    push_unique_subject_token(&mut subjects, token.id_token.as_str());
-    subjects
+fn api_key_exchange_subject_token(token: &Token) -> Option<String> {
+    // `/oauth/token` uses the token-exchange grant and expects the OAuth ID
+    // token as its subject.  An access token is the bearer fallback for the
+    // upstream request; sending it to this endpoint produces misleading
+    // "Invalid ID token" / audience errors even when the account is usable.
+    let id_token = token.id_token.trim();
+    (!id_token.is_empty()).then(|| id_token.to_string())
 }
 
-fn push_unique_subject_token(subjects: &mut Vec<String>, candidate: &str) {
-    let value = candidate.trim();
-    if value.is_empty() || subjects.iter().any(|existing| existing == value) {
-        return;
-    }
-    subjects.push(value.to_string());
+pub(crate) fn api_key_exchange_client_id(token: &Token, fallback_client_id: &str) -> String {
+    // The exchange subject is the ID token, so prefer its client_id claim.  A
+    // separately issued access token can carry a different audience/client
+    // claim and must not override the ID-token exchange client.
+    extract_client_id_claim(&token.id_token)
+        .or_else(|| extract_client_id_claim(&token.access_token))
+        .or_else(|| {
+            let fallback = fallback_client_id.trim();
+            (!fallback.is_empty()).then(|| fallback.to_string())
+        })
+        .unwrap_or_else(|| DEFAULT_CLIENT_ID.to_string())
 }
 
 /// 函数 `fallback_to_access_token`
@@ -256,7 +255,7 @@ pub(super) fn resolve_openai_bearer_token(
     }
 
     let fallback_client_id = super::runtime_config::token_exchange_client_id();
-    let client_id = crate::usage_token_refresh::token_refresh_client_id(token, &fallback_client_id);
+    let client_id = api_key_exchange_client_id(token, &fallback_client_id);
     let issuer_env = super::runtime_config::token_exchange_default_issuer();
     let issuer = if account.issuer.trim().is_empty() {
         issuer_env
@@ -305,10 +304,7 @@ pub(super) fn resolve_openai_bearer_token(
                         let _ = storage.insert_token(token);
 
                         if !token.id_token.trim().is_empty() {
-                            let refreshed_client_id =
-                                crate::usage_token_refresh::token_refresh_client_id(
-                                    token, &client_id,
-                                );
+                            let refreshed_client_id = api_key_exchange_client_id(token, &client_id);
                             if let Ok(exchanged) = exchange_and_persist_api_key_access_token(
                                 storage,
                                 token,

crates/service/src/gateway/mod.rs

@@ -436,6 +436,7 @@ pub(crate) use selection::{
 };
 #[cfg(test)]
 use token_exchange::account_token_exchange_lock;
+pub(crate) use token_exchange::api_key_exchange_client_id;
 use token_exchange::resolve_openai_bearer_token;
 use upstream::proxy::proxy_validated_request;
 

crates/service/src/gateway/upstream/attempt_flow/transport.rs

@@ -1481,11 +1481,7 @@ fn is_websocket_upstream_terminal_text(text: &str) -> bool {
             .unwrap_or_default()
             .to_ascii_lowercase()
             .as_str(),
-        "response.completed"
-            | "response.done"
-            | "response.failed"
-            | "response.incomplete"
-            | "error"
+        "response.completed" | "response.failed" | "response.incomplete" | "error"
     )
 }
 
@@ -1521,7 +1517,7 @@ fn is_websocket_upstream_connection_limit_text(text: &str) -> bool {
 }
 
 fn is_websocket_upstream_transport_healthy_terminal_text(text: &str) -> bool {
-    is_websocket_upstream_terminal_text(text) && !is_websocket_upstream_connection_limit_text(text)
+    is_websocket_upstream_completed_text(text) && !is_websocket_upstream_connection_limit_text(text)
 }
 
 fn websocket_upstream_sse_event(text: &str) -> String {
@@ -1533,7 +1529,6 @@ fn websocket_upstream_sse_event(text: &str) -> String {
     }
 }
 
-#[cfg(test)]
 fn is_websocket_upstream_completed_text(text: &str) -> bool {
     serde_json::from_str::<serde_json::Value>(text)
         .ok()

crates/service/src/gateway/upstream/attempt_flow/transport_tests.rs

@@ -807,7 +807,7 @@ fn websocket_upstream_terminal_detection_parses_json_type() {
     assert!(super::is_websocket_upstream_terminal_text(
         r#"{"type":"response.completed"}"#
     ));
-    assert!(super::is_websocket_upstream_terminal_text(
+    assert!(!super::is_websocket_upstream_terminal_text(
         r#"{"type":"response.done"}"#
     ));
     assert!(super::is_websocket_upstream_terminal_text(
@@ -991,7 +991,7 @@ fn send_websocket_upstream_request_builds_valid_handshake_and_stops_on_completed
 }
 
 #[test]
-fn send_websocket_upstream_request_does_not_cooldown_after_application_failure() {
+fn send_websocket_upstream_request_does_not_mark_recovery_completed_after_application_failure() {
     let _env_lock = crate::test_env_guard();
     let _reload_guard = RuntimeConfigReloadGuard;
     let _proxy_guard = EnvGuard::set("CODEXMANAGER_UPSTREAM_PROXY_URL", "");
@@ -1016,6 +1016,11 @@ fn send_websocket_upstream_request_does_not_cooldown_after_application_failure()
 
     assert!(
         super::is_websocket_upstream_transport_healthy_terminal_text(
+            r#"{"type":"response.completed"}"#
+        )
+    );
+    assert!(
+        !super::is_websocket_upstream_transport_healthy_terminal_text(
             r#"{"type":"response.failed"}"#
         )
     );