From ef57b880662d6e4b5b6f55ebd678deb7855c01fc Mon Sep 17 00:00:00 2001 From: zyycn Date: Sat, 25 Jul 2026 03:26:00 +0800 Subject: [PATCH] fix(openai): preserve Agent Identity CPR round trips Co-authored-by: Codex --- .../providers/openai/src/credential/admin.rs | 164 +++++++++++++----- .../openai/tests/credential/admin.rs | 125 ++++++++++++- 2 files changed, 238 insertions(+), 51 deletions(-) diff --git a/backend/crates/providers/openai/src/credential/admin.rs b/backend/crates/providers/openai/src/credential/admin.rs index 380540dc4..e5cb51f9a 100644 --- a/backend/crates/providers/openai/src/credential/admin.rs +++ b/backend/crates/providers/openai/src/credential/admin.rs @@ -215,7 +215,7 @@ impl fmt::Debug for ExportManagedCodexCredential { } } -/// CPR canonical 账号导出文档;只允许显式序列化,Debug 永不输出 token。 +/// CPR canonical 账号导出文档;只允许显式序列化,Debug 永不输出 credential secret。 #[derive(Serialize)] #[serde(rename_all = "camelCase")] pub struct CodexCprExportDocument { @@ -252,21 +252,47 @@ impl fmt::Debug for CodexCprExportDocument { #[derive(Serialize)] #[serde(rename_all = "camelCase")] -struct CodexCprExportAccount { +struct CodexCprExportCommon { id: String, email: Option, account_id: Option, user_id: Option, label: Option, plan_type: Option, - token: String, - refresh_token: Option, - access_token_expires_at: Option, status: &'static str, added_at: String, updated_at: String, } +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct CodexCprOAuthExportAccount { + #[serde(flatten)] + common: CodexCprExportCommon, + token: String, + refresh_token: Option, + access_token_expires_at: Option, +} + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct CodexCprAgentIdentityExportAccount { + #[serde(flatten)] + common: CodexCprExportCommon, + auth_mode: &'static str, + agent_runtime_id: String, + agent_private_key: String, + #[serde(skip_serializing_if = "Option::is_none")] + task_id: Option, +} + +#[derive(Serialize)] +#[serde(untagged)] +enum CodexCprExportAccount { + OAuth(CodexCprOAuthExportAccount), + AgentIdentity(CodexCprAgentIdentityExportAccount), +} + /// App 已从 Store 读取的当前账号、revision 与明文 Provider JSON。 pub struct RotateManagedCodexCredential { pub current: LoadedCredential, @@ -444,31 +470,51 @@ impl CodexCredentialAdmin { } let data = CodexCredentialCodec::decode_complete(&item.current.credential) .map_err(|_| CodexCredentialAdminError::InvalidCredential)?; - let CodexCredentialData::OAuth(data) = data else { - return Err(CodexCredentialAdminError::InvalidCredential); - }; - if account.authentication_kind() != CODEX_AUTHENTICATION_KIND_OAUTH - || account.has_refresh_token() != data.refresh_token.is_some() - { - return Err(CodexCredentialAdminError::InvalidCredential); - } - accounts.push(CodexCprExportAccount { + let common = CodexCprExportCommon { id: account.id().as_str().to_owned(), email: account.email().map(str::to_owned), account_id: account.upstream_account_id().map(str::to_owned), user_id: Some(account.upstream_user_id().to_owned()), label: Some(account.name().to_owned()), plan_type: account.plan_type().map(str::to_owned), - token: data.access_token, - refresh_token: data.refresh_token, - access_token_expires_at: account - .access_token_expires_at() - .map(DateTime::::from) - .map(|value| value.to_rfc3339()), status: cpr_status(&account), added_at: china_rfc3339(item.added_at), updated_at: china_rfc3339(item.updated_at), - }); + }; + let exported = match data { + CodexCredentialData::OAuth(data) => { + if account.authentication_kind() != CODEX_AUTHENTICATION_KIND_OAUTH + || account.has_refresh_token() != data.refresh_token.is_some() + { + return Err(CodexCredentialAdminError::InvalidCredential); + } + CodexCprExportAccount::OAuth(CodexCprOAuthExportAccount { + common, + token: data.access_token, + refresh_token: data.refresh_token, + access_token_expires_at: account + .access_token_expires_at() + .map(DateTime::::from) + .map(|value| value.to_rfc3339()), + }) + } + CodexCredentialData::AgentIdentity(data) => { + if account.authentication_kind() != CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY + || account.has_refresh_token() + || account.access_token_expires_at().is_some() + { + return Err(CodexCredentialAdminError::InvalidCredential); + } + CodexCprExportAccount::AgentIdentity(CodexCprAgentIdentityExportAccount { + common, + auth_mode: "agentIdentity", + agent_runtime_id: data.agent_runtime_id, + agent_private_key: data.agent_private_key, + task_id: data.task_id, + }) + } + }; + accounts.push(exported); } Ok(CodexCprExportDocument { source_format: "cpr", @@ -1034,7 +1080,25 @@ fn parse_import_document( let values = import_account_values(payload)?; let mut accounts = Vec::new(); for value in values { - let parsed = match shape { + // CPR 导出允许 OAuth 与 Agent Identity 共存;按账号选择 Provider + // 解析器,避免一个 Agent Identity 条目把整份文档误判成单一格式。 + let account_shape = if looks_like_agent_identity_account(value) { + ImportDocumentShape::AgentIdentity + } else if shape == ImportDocumentShape::AgentIdentity { + if value.get("credentials").is_some() || value.get("tokens").is_some() { + ImportDocumentShape::CredentialBundle + } else if value + .as_object() + .is_some_and(auth_document_provider_is_openai) + { + ImportDocumentShape::AuthDocument + } else { + ImportDocumentShape::Native + } + } else { + shape + }; + let parsed = match account_shape { ImportDocumentShape::Native => Some(parse_native_account(value)?), ImportDocumentShape::AgentIdentity => Some(parse_agent_identity_account(value)?), ImportDocumentShape::CredentialBundle => parse_credential_bundle_account(value)?, @@ -1076,26 +1140,37 @@ fn parse_agent_identity_account( .get("agent_identity") .or_else(|| value.get("credentials")) .unwrap_or(value); - let auth_mode = - first_string(value, &["auth_mode"]).or_else(|| first_string(identity, &["auth_mode"])); + let auth_mode = first_string(value, &["auth_mode", "authMode"]) + .or_else(|| first_string(identity, &["auth_mode", "authMode"])); if auth_mode.as_deref() != Some("agentIdentity") { return Err(CodexCredentialAdminError::InvalidCredential); } - let runtime_id = first_string(identity, &["agent_runtime_id"]) + let runtime_id = first_string(identity, &["agent_runtime_id", "agentRuntimeId"]) .ok_or(CodexCredentialAdminError::InvalidCredential)?; - let private_key = first_string(identity, &["agent_private_key"]) + let private_key = first_string(identity, &["agent_private_key", "agentPrivateKey"]) .ok_or(CodexCredentialAdminError::InvalidCredential)?; Ok(ParsedCodexImportAccount { id: first_string(value, &["id"]), name: first_string(value, &["name", "label"]), email: first_string(identity, &["email"]), - plan_type: first_string(identity, &["plan_type"]), - chatgpt_account_id: first_string(identity, &["chatgpt_account_id", "account_id"]), - chatgpt_user_id: first_string(identity, &["chatgpt_user_id"]), + plan_type: first_string(identity, &["plan_type", "planType"]), + chatgpt_account_id: first_string( + identity, + &[ + "chatgpt_account_id", + "chatgptAccountId", + "account_id", + "accountId", + ], + ), + chatgpt_user_id: first_string( + identity, + &["chatgpt_user_id", "chatgptUserId", "user_id", "userId"], + ), authentication: ParsedCodexAuthentication::AgentIdentity { runtime_id, private_key, - task_id: first_string(identity, &["task_id"]), + task_id: first_string(identity, &["task_id", "taskId"]), }, status: first_string(value, &["status"]), }) @@ -1379,24 +1454,23 @@ fn looks_like_credential_bundle(value: &Value) -> bool { fn looks_like_agent_identity_document(value: &Value) -> bool { import_account_values(value).is_ok_and(|accounts| { - accounts.iter().any(|account| { - account - .get("auth_mode") - .and_then(Value::as_str) - .is_some_and(|mode| mode == "agentIdentity") - || account - .get("agent_identity") - .and_then(|identity| identity.get("agent_runtime_id")) - .is_some() - || account - .get("credentials") - .and_then(|credentials| credentials.get("auth_mode")) - .and_then(Value::as_str) - .is_some_and(|mode| mode == "agentIdentity") - }) + accounts + .iter() + .any(|account| looks_like_agent_identity_account(account)) }) } +fn looks_like_agent_identity_account(value: &Value) -> bool { + let identity = value + .get("agent_identity") + .or_else(|| value.get("credentials")) + .unwrap_or(value); + first_string(value, &["auth_mode", "authMode"]) + .or_else(|| first_string(identity, &["auth_mode", "authMode"])) + .is_some_and(|mode| mode == "agentIdentity") + || first_string(identity, &["agent_runtime_id", "agentRuntimeId"]).is_some() +} + fn looks_like_auth_document(value: &Value) -> bool { import_account_values(value).is_ok_and(|accounts| { accounts.iter().any(|account| { diff --git a/backend/crates/providers/openai/tests/credential/admin.rs b/backend/crates/providers/openai/tests/credential/admin.rs index b7307c4e6..3c9029b21 100644 --- a/backend/crates/providers/openai/tests/credential/admin.rs +++ b/backend/crates/providers/openai/tests/credential/admin.rs @@ -3,23 +3,31 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use async_trait::async_trait; +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD; use chrono::{TimeZone, Utc}; +use ed25519_dalek::SigningKey; +use ed25519_dalek::pkcs8::EncodePrivateKey as _; use futures::future::BoxFuture; use gateway_core::engine::credential::{ - CredentialRevision, LoadedCredential, PlaintextCredential, ProviderAccountId, + AccountAvailability, CredentialRevision, LoadedCredential, PlaintextCredential, + ProviderAccount, ProviderAccountId, }; use gateway_core::provider_ports::{ ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshLeaseRequest, ProviderStoreError, }; +use gateway_core::routing::ProviderKind; use provider_openai::OFFICIAL_CODEX_BASE_URL; use provider_openai::credential::token_client::{RefreshFailure, TokenPair, TokenRefresher}; use provider_openai::credential::{ - CodexAccountIdentityService, CodexAccountIdentityVerifier, CodexCredentialAdmin, - CodexCredentialAdminError, CodexCredentialAdminService, CodexCredentialCodec, - CodexIdentityExpectation, CodexIdentityVerification, CodexIdentityVerificationError, - CodexJwtIdentityVerifier, ExportManagedCodexCredential, ImportCodexOAuthCredential, - ReqwestCodexAuthenticatedAccountSource, ReqwestOpenAiJwksSource, RotateManagedCodexCredential, + CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY, CodexAccountIdentityService, + CodexAccountIdentityVerifier, CodexAgentIdentityAuthMode, CodexAgentIdentityCredentialData, + CodexCredentialAdmin, CodexCredentialAdminError, CodexCredentialAdminService, + CodexCredentialCodec, CodexIdentityExpectation, CodexIdentityVerification, + CodexIdentityVerificationError, CodexJwtIdentityVerifier, ExportManagedCodexCredential, + ImportCodexOAuthCredential, ReqwestCodexAuthenticatedAccountSource, ReqwestOpenAiJwksSource, + RotateManagedCodexCredential, }; use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; use secrecy::{ExposeSecret, SecretString}; @@ -137,6 +145,60 @@ fn export_item(id: &str, token: &str) -> ExportManagedCodexCredential { } } +fn agent_export_item(id: &str) -> (ExportManagedCodexCredential, String) { + let signing_key = SigningKey::from_bytes(&[7; 32]); + let private_key = STANDARD.encode( + signing_key + .to_pkcs8_der() + .expect("encode Agent Identity key") + .as_bytes(), + ); + let credential = + CodexCredentialCodec::encode_agent_identity(CodexAgentIdentityCredentialData { + schema_version: 1, + auth_mode: CodexAgentIdentityAuthMode::AgentIdentity, + installation_id: "00000000-0000-4000-8000-000000000007".to_owned(), + agent_runtime_id: "runtime-export".to_owned(), + agent_private_key: private_key.clone(), + task_id: Some("task-export".to_owned()), + cookies: Vec::new(), + }) + .expect("encode Agent Identity credential"); + let account = ProviderAccount::new( + ProviderAccountId::new(id.to_owned()).expect("account id"), + ProviderKind::new("openai").expect("provider"), + format!("Agent Identity {id}"), + "agent-user-export".to_owned(), + CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY.to_owned(), + CredentialRevision::new(1).expect("revision"), + None, + ) + .with_profile( + Some("agent-export@example.com".to_owned()), + Some("agent-account-export".to_owned()), + Some("pro".to_owned()), + ) + .with_runtime_state(true, AccountAvailability::Ready, None) + .with_refresh_schedule(false, None); + ( + ExportManagedCodexCredential { + current: LoadedCredential { + account, + credential, + }, + added_at: Utc + .with_ymd_and_hms(2026, 7, 18, 2, 47, 1) + .single() + .expect("added at"), + updated_at: Utc + .with_ymd_and_hms(2026, 7, 19, 3, 0, 0) + .single() + .expect("updated at"), + }, + private_key, + ) +} + #[test] fn cpr_export_matches_the_canonical_real_document_shape() { let document = CodexCredentialAdmin @@ -211,6 +273,57 @@ fn cpr_export_batch_validation_and_debug_are_secret_safe() { assert!(debug.contains("account_count: 1")); } +#[test] +fn cpr_export_preserves_agent_identity_material_without_oauth_fields() { + let (item, private_key) = agent_export_item("acct_agent_export"); + let value = CodexCredentialAdmin + .format_cpr_export(vec![item]) + .expect("Agent Identity export") + .into_json() + .expect("serialize Agent Identity export"); + let account = &value["accounts"][0]; + + assert_eq!(account["authMode"], "agentIdentity"); + assert_eq!(account["agentRuntimeId"], "runtime-export"); + assert!( + account["agentPrivateKey"] + .as_str() + .is_some_and(|value| value == private_key) + ); + assert_eq!(account["taskId"], "task-export"); + assert!(account.get("token").is_none()); + assert!(account.get("refreshToken").is_none()); + assert!(account.get("accessTokenExpiresAt").is_none()); +} + +#[tokio::test] +async fn cpr_export_and_import_support_mixed_oauth_and_agent_identity_accounts() { + let (agent, _) = agent_export_item("acct_mixed_agent"); + let value = CodexCredentialAdmin + .format_cpr_export(vec![ + export_item("acct_mixed_oauth", "token-acct_mixed_oauth"), + agent, + ]) + .expect("mixed export") + .into_json() + .expect("serialize mixed export"); + + let prepared = import_service(unused_import_refresher()) + .prepare_import_document(value) + .await + .expect("mixed CPR import"); + assert_eq!(prepared.accounts().len(), 2); + assert!(prepared.accounts().iter().any(|account| { + account.account.authentication_kind() == CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY + })); + assert!( + prepared + .accounts() + .iter() + .any(|account| { account.account.authentication_kind() == "oauth" }) + ); +} + pub(super) struct ManualRefresher { outcome: Mutex>>, seen: Mutex>,