feat(codex): preserve OAuth login state during third-party provider switching

Codex provider switches now only write config.toml for third-party providers,
injecting the API key as experimental_bearer_token. The user's auth.json
(ChatGPT OAuth tokens) is preserved. Official providers with login material
still write auth.json normally. Backfill restores bearer tokens into stored
provider auth.OPENAI_API_KEY to maintain canonical shape.
This commit is contained in:
Jason
2026-05-25 17:46:23 +08:00
parent 88ba908bc4
commit 95f2dd4126
17 changed files with 1664 additions and 154 deletions
+447 -14
View File
@@ -1,14 +1,12 @@
// unused imports removed
use std::path::PathBuf;
use std::path::{Path, PathBuf};
use crate::config::{
atomic_write, delete_file, get_home_dir, sanitize_provider_name, write_json_file,
write_text_file,
atomic_write, delete_file, get_home_dir, read_json_file, sanitize_provider_name,
write_json_file, write_text_file,
};
use crate::error::AppError;
use serde_json::{json, Value};
use std::fs;
use std::path::Path;
use std::process::Command;
use toml_edit::DocumentMut;
@@ -393,20 +391,123 @@ pub fn write_codex_live_atomic_with_stable_provider(
) -> Result<(), AppError> {
match config_text_opt {
Some(config_text) => {
let mut settings = serde_json::Map::new();
settings.insert("config".to_string(), Value::String(config_text.to_string()));
let mut settings = Value::Object(settings);
normalize_codex_settings_config_model_provider(&mut settings)?;
let config_text = settings
.get("config")
.and_then(|value| value.as_str())
.unwrap_or(config_text);
write_codex_live_atomic(auth, Some(config_text))
let config_text = normalize_codex_config_for_live_provider(config_text)?;
write_codex_live_atomic(auth, Some(&config_text))
}
None => write_codex_live_atomic(auth, None),
}
}
fn normalize_codex_config_for_live_provider(config_text: &str) -> Result<String, AppError> {
let mut settings = serde_json::Map::new();
settings.insert("config".to_string(), Value::String(config_text.to_string()));
let mut settings = Value::Object(settings);
normalize_codex_settings_config_model_provider(&mut settings)?;
Ok(settings
.get("config")
.and_then(|value| value.as_str())
.unwrap_or(config_text)
.to_string())
}
/// Write only Codex `config.toml` for provider switching.
///
/// Codex login state lives in `auth.json`; provider routing, endpoint, model,
/// and provider-scoped bearer tokens live in `config.toml`. Provider switches
/// should not overwrite the user's ChatGPT login cache.
pub fn write_codex_live_config_atomic_with_stable_provider(
config_text_opt: Option<&str>,
) -> Result<(), AppError> {
let config_path = get_codex_config_path();
let cfg_text = match config_text_opt {
Some(config_text) => normalize_codex_config_for_live_provider(config_text)?,
None => String::new(),
};
if !cfg_text.trim().is_empty() {
toml::from_str::<toml::Table>(&cfg_text).map_err(|e| AppError::toml(&config_path, e))?;
}
write_text_file(&config_path, &cfg_text)
}
pub fn extract_codex_auth_api_key(auth: &Value) -> Option<String> {
auth.get("OPENAI_API_KEY")
.and_then(|value| value.as_str())
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
}
pub fn extract_codex_api_key(auth: Option<&Value>, config_text: Option<&str>) -> Option<String> {
auth.and_then(extract_codex_auth_api_key)
.or_else(|| config_text.and_then(extract_codex_experimental_bearer_token))
}
pub fn codex_auth_has_login_material(auth: &Value) -> bool {
let Some(obj) = auth.as_object() else {
return false;
};
obj.iter().any(|(key, value)| {
if key == "auth_mode" {
return false;
}
if key == "OPENAI_API_KEY" {
return value
.as_str()
.map(str::trim)
.is_some_and(|token| !token.is_empty());
}
match value {
Value::Null => false,
Value::String(text) => !text.trim().is_empty(),
Value::Array(items) => !items.is_empty(),
Value::Object(map) => !map.is_empty(),
_ => true,
}
})
}
pub fn codex_auth_has_oauth_login_material(auth: &Value) -> bool {
let Some(obj) = auth.as_object() else {
return false;
};
obj.iter().any(|(key, value)| {
if key == "auth_mode" || key == "OPENAI_API_KEY" {
return false;
}
match value {
Value::Null => false,
Value::String(text) => !text.trim().is_empty(),
Value::Array(items) => !items.is_empty(),
Value::Object(map) => !map.is_empty(),
_ => true,
}
})
}
pub fn should_restore_codex_provider_token_for_backfill(
category: Option<&str>,
template_settings: &Value,
) -> bool {
if category == Some("official") {
return false;
}
let Some(auth) = template_settings.get("auth") else {
return true;
};
let has_provider_api_key = extract_codex_auth_api_key(auth).is_some();
let has_oauth_login = codex_auth_has_oauth_login_material(auth);
!(has_oauth_login && !has_provider_api_key)
}
fn parse_codex_positive_u64(value: Option<&Value>) -> Option<u64> {
match value {
Some(Value::Number(n)) => n.as_u64().filter(|v| *v > 0),
@@ -787,6 +888,235 @@ pub fn write_codex_live_with_catalog(
write_codex_live_atomic_with_stable_provider(auth, prepared_config.as_deref())
}
pub fn write_codex_provider_live_with_catalog(
settings: &Value,
category: Option<&str>,
auth: &Value,
config_text: Option<&str>,
) -> Result<(), AppError> {
let prepared_config = config_text
.map(|text| prepare_codex_config_text_with_model_catalog(settings, text))
.transpose()?;
write_codex_live_for_provider(category, auth, prepared_config.as_deref())
}
/// Extract a provider-scoped `experimental_bearer_token` from Codex `config.toml`.
///
/// Mobile compat: third-party providers may store the API key inside
/// `[model_providers.<id>].experimental_bearer_token` while keeping the
/// user's ChatGPT login cache intact in `auth.json`. Falls back to the
/// top-level `experimental_bearer_token` when no active model provider is set.
pub fn extract_codex_experimental_bearer_token(config_text: &str) -> Option<String> {
if !config_text.contains("experimental_bearer_token") {
return None;
}
let doc = config_text.parse::<DocumentMut>().ok()?;
let provider_id = active_codex_model_provider_id(&doc);
let top_level_token = || {
doc.get("experimental_bearer_token")
.and_then(|item| item.as_str())
};
let token = match provider_id.as_deref() {
Some(id) if is_custom_codex_model_provider_id(id) => doc
.get("model_providers")
.and_then(|item| item.as_table())
.and_then(|table| table.get(id))
.and_then(|item| item.as_table())
.and_then(|table| table.get("experimental_bearer_token"))
.and_then(|item| item.as_str())
.or_else(top_level_token),
Some(_) => top_level_token(),
None => top_level_token(),
};
token
.map(str::trim)
.filter(|token| !token.is_empty())
.map(str::to_string)
}
fn set_codex_experimental_bearer_token(config_text: &str, token: &str) -> Result<String, AppError> {
if config_text.trim().is_empty() {
return Err(AppError::localized(
"provider.codex.config.missing",
"Codex 第三方供应商缺少 config.toml 配置,无法写入 bearer token",
"Codex third-party provider is missing config.toml, cannot write bearer token",
));
}
let mut doc = config_text
.parse::<DocumentMut>()
.map_err(|e| AppError::Message(format!("Invalid Codex config.toml: {e}")))?;
let Some(provider_id) = active_codex_model_provider_id(&doc) else {
doc["experimental_bearer_token"] = toml_edit::value(token);
return Ok(doc.to_string());
};
if !is_custom_codex_model_provider_id(&provider_id) {
// Reserved Codex provider IDs are owned by the CLI. Keep third-party
// bearer tokens at the top level so we do not shadow built-in tables.
doc["experimental_bearer_token"] = toml_edit::value(token);
return Ok(doc.to_string());
}
if let Some(model_providers) = doc
.get_mut("model_providers")
.and_then(|item| item.as_table_mut())
{
if let Some(provider_table) = model_providers
.get_mut(provider_id.as_str())
.and_then(|item| item.as_table_mut())
{
provider_table["experimental_bearer_token"] = toml_edit::value(token);
return Ok(doc.to_string());
}
}
doc["experimental_bearer_token"] = toml_edit::value(token);
Ok(doc.to_string())
}
fn remove_codex_experimental_bearer_token(config_text: &str) -> Result<String, AppError> {
if config_text.trim().is_empty() || !config_text.contains("experimental_bearer_token") {
return Ok(config_text.to_string());
}
let mut doc = config_text
.parse::<DocumentMut>()
.map_err(|e| AppError::Message(format!("Invalid Codex config.toml: {e}")))?;
if let Some(provider_id) = active_codex_model_provider_id(&doc) {
if let Some(provider_table) = doc
.get_mut("model_providers")
.and_then(|item| item.as_table_mut())
.and_then(|table| table.get_mut(provider_id.as_str()))
.and_then(|item| item.as_table_mut())
{
provider_table.remove("experimental_bearer_token");
}
}
doc.as_table_mut().remove("experimental_bearer_token");
Ok(doc.to_string())
}
/// Read the current Codex live settings as a `{ auth, config }` object.
///
/// Missing `auth.json` collapses to `{}` so a config-only third-party install
/// is still importable; both files empty is treated as "no live install".
pub fn read_codex_live_settings() -> Result<Value, AppError> {
let auth_path = get_codex_auth_path();
let auth_present = auth_path.exists();
let auth: Value = if auth_present {
read_json_file(&auth_path)?
} else {
json!({})
};
let cfg_text = read_and_validate_codex_config_text()?;
if !auth_present && cfg_text.trim().is_empty() {
return Err(AppError::localized(
"codex.live.missing",
"Codex 配置文件不存在",
"Codex configuration is missing",
));
}
Ok(json!({ "auth": auth, "config": cfg_text }))
}
/// Route a Codex live write between full auth+config or config-only.
///
/// Official providers with usable login material own `auth.json`; everyone
/// else only touches `config.toml` so the user's ChatGPT login cache survives
/// third-party switches.
pub fn write_codex_live_for_provider(
category: Option<&str>,
auth: &Value,
config_text: Option<&str>,
) -> Result<(), AppError> {
if category == Some("official") && codex_auth_has_login_material(auth) {
write_codex_live_atomic_with_stable_provider(auth, config_text)
} else {
let live_config = prepare_codex_provider_live_config(auth, config_text.unwrap_or(""))?;
write_codex_live_config_atomic_with_stable_provider(Some(&live_config))
}
}
/// Build the live Codex config for provider switching.
///
/// The stored provider keeps its API key in `auth.OPENAI_API_KEY`. Live Codex
/// requests can use a provider-scoped `experimental_bearer_token`, so switching
/// providers only needs to update `config.toml`; `auth.json` stays as the user's
/// long-lived ChatGPT login cache.
pub fn prepare_codex_provider_live_config(
auth: &Value,
config_text: &str,
) -> Result<String, AppError> {
let token = extract_codex_auth_api_key(auth)
.or_else(|| extract_codex_experimental_bearer_token(config_text));
Ok(match token {
Some(token) => set_codex_experimental_bearer_token(config_text, &token)?,
None => config_text.to_string(),
})
}
/// During DB backfill, lift a live `experimental_bearer_token` back into
/// `auth.OPENAI_API_KEY` so the stored provider keeps its canonical shape
/// and generated live tokens don't leak into stored provider TOML.
///
/// Only intervenes when the live config actually carries a bearer token —
/// otherwise the function is a no-op so the caller's normal backfill path
/// (which keeps live `auth` as the authoritative source) is unaffected.
pub fn restore_codex_provider_token_for_backfill(
settings: &mut Value,
template_settings: &Value,
) -> Result<(), AppError> {
let Some(config_text) = settings
.get("config")
.and_then(|value| value.as_str())
.map(str::to_string)
else {
return Ok(());
};
let Some(token) = extract_codex_experimental_bearer_token(&config_text) else {
return Ok(());
};
let cleaned_config = remove_codex_experimental_bearer_token(&config_text)?;
if let Some(obj) = settings.as_object_mut() {
obj.insert("config".to_string(), Value::String(cleaned_config));
let mut auth = template_settings
.get("auth")
.filter(|value| value.is_object())
.cloned()
.unwrap_or_else(|| Value::Object(serde_json::Map::new()));
if let Some(auth_obj) = auth.as_object_mut() {
auth_obj.insert("OPENAI_API_KEY".to_string(), Value::String(token));
}
obj.insert("auth".to_string(), auth);
}
Ok(())
}
pub fn restore_codex_settings_for_backfill(
settings: &mut Value,
template_settings: &Value,
restore_provider_token: bool,
) -> Result<(), AppError> {
restore_codex_settings_config_model_provider_for_backfill(settings, template_settings)?;
if restore_provider_token {
restore_codex_provider_token_for_backfill(settings, template_settings)?;
}
Ok(())
}
/// Update a field in Codex config.toml using toml_edit (syntax-preserving).
///
/// Supported fields:
@@ -905,6 +1235,7 @@ pub fn remove_codex_toml_base_url_if(toml_str: &str, predicate: impl Fn(&str) ->
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn normalize_live_config_uses_custom_for_third_party_model_provider_id() {
@@ -984,6 +1315,108 @@ wire_api = "responses"
assert_eq!(result, "");
}
#[test]
fn prepare_provider_live_config_rejects_key_without_config() {
let err = prepare_codex_provider_live_config(&json!({"OPENAI_API_KEY": "sk-test"}), "")
.expect_err("empty config with API key should not truncate live config");
assert!(
err.to_string().contains("config.toml"),
"error should explain missing config.toml, got: {err}"
);
}
#[test]
fn prepare_provider_live_config_uses_top_level_token_for_reserved_provider() {
let input = r#"model_provider = "openai"
model = "gpt-5"
"#;
let output =
prepare_codex_provider_live_config(&json!({"OPENAI_API_KEY": "sk-test"}), input)
.expect("prepare live config");
let parsed: toml::Value = toml::from_str(&output).expect("parse output");
assert_eq!(
parsed
.get("experimental_bearer_token")
.and_then(|v| v.as_str()),
Some("sk-test")
);
assert!(
parsed.get("model_providers").is_none(),
"reserved provider tables should not be synthesized"
);
}
#[test]
fn extract_bearer_uses_top_level_token_for_reserved_provider() {
let input = r#"model_provider = "openai"
experimental_bearer_token = "top-level-key"
[model_providers.openai]
experimental_bearer_token = "stale-table-key"
"#;
assert_eq!(
extract_codex_experimental_bearer_token(input).as_deref(),
Some("top-level-key")
);
}
#[test]
fn should_not_restore_provider_token_for_oauth_only_template() {
let oauth_template = json!({
"auth": {
"auth_mode": "chatgpt",
"tokens": {
"access_token": "oauth-access"
}
}
});
let api_key_template = json!({
"auth": {
"OPENAI_API_KEY": "sk-test"
}
});
assert!(
!should_restore_codex_provider_token_for_backfill(Some("custom"), &oauth_template),
"OAuth-only templates should not backfill bearer tokens into OPENAI_API_KEY"
);
assert!(
should_restore_codex_provider_token_for_backfill(Some("custom"), &api_key_template),
"custom API-key providers should still restore provider bearer tokens"
);
assert!(
!should_restore_codex_provider_token_for_backfill(Some("official"), &api_key_template),
"official providers should never restore third-party bearer tokens"
);
}
#[test]
fn prepare_provider_live_config_does_not_create_incomplete_provider_table() {
let input = r#"model_provider = "vendor_x"
model = "gpt-5"
"#;
let output =
prepare_codex_provider_live_config(&json!({"OPENAI_API_KEY": "sk-test"}), input)
.expect("prepare live config");
let parsed: toml::Value = toml::from_str(&output).expect("parse output");
assert_eq!(
parsed
.get("experimental_bearer_token")
.and_then(|v| v.as_str()),
Some("sk-test")
);
assert!(
parsed.get("model_providers").is_none(),
"missing provider tables should not be synthesized without endpoint fields"
);
}
#[test]
fn normalize_live_config_rewrites_matching_profile_model_provider_refs() {
let target = r#"model_provider = "vendor_alpha"