From cf801bad4dbe34ad6aa75e0fa775e13359fa91cc Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Wed, 18 Mar 2026 12:21:10 -0700 Subject: [PATCH] Extract provider and token modules into codex-auth Move the foundational provider and token modules into codex-auth while keeping codex-core as the facade. Also move the corresponding unit tests and record the 3-PR migration checkpoints. Co-authored-by: Codex --- codex-rs/.codex/codex-auth-migration-plan.md | 171 +++++++++ codex-rs/Cargo.lock | 19 + codex-rs/Cargo.toml | 2 + codex-rs/codex-auth/BUILD.bazel | 6 + codex-rs/codex-auth/Cargo.toml | 24 ++ codex-rs/codex-auth/src/error.rs | 15 + codex-rs/codex-auth/src/lib.rs | 19 + .../src/model_provider_info_tests.rs | 2 +- codex-rs/codex-auth/src/provider.rs | 286 ++++++++++++++ codex-rs/codex-auth/src/token_data.rs | 163 ++++++++ .../src/token_data_tests.rs | 3 +- codex-rs/core/Cargo.toml | 1 + codex-rs/core/src/client.rs | 2 +- codex-rs/core/src/error.rs | 27 +- codex-rs/core/src/model_provider_info.rs | 357 +----------------- codex-rs/core/src/models_manager/manager.rs | 4 +- codex-rs/core/src/realtime_conversation.rs | 3 +- codex-rs/core/src/token_data.rs | 180 +-------- 18 files changed, 724 insertions(+), 560 deletions(-) create mode 100644 codex-rs/.codex/codex-auth-migration-plan.md create mode 100644 codex-rs/codex-auth/BUILD.bazel create mode 100644 codex-rs/codex-auth/Cargo.toml create mode 100644 codex-rs/codex-auth/src/error.rs create mode 100644 codex-rs/codex-auth/src/lib.rs rename codex-rs/{core => codex-auth}/src/model_provider_info_tests.rs (99%) create mode 100644 codex-rs/codex-auth/src/provider.rs create mode 100644 codex-rs/codex-auth/src/token_data.rs rename codex-rs/{core => codex-auth}/src/token_data_tests.rs (98%) diff --git a/codex-rs/.codex/codex-auth-migration-plan.md b/codex-rs/.codex/codex-auth-migration-plan.md new file mode 100644 index 0000000000..552ae54f83 --- /dev/null +++ b/codex-rs/.codex/codex-auth-migration-plan.md @@ -0,0 +1,171 @@ +# Codex Auth Migration Plan + +## Goal + +Move auth and model-provider ownership into a new `codex-auth` crate while keeping `codex-core` as the thin public surface for the rest of the workspace. + +Non-`core` crates should not end up in a state where they need to import `codex-auth` directly. + +## Invariants + +- `codex-auth` is the implementation owner for auth and provider concepts. +- `codex-core` is the public facade for normal consumers. +- Normal consumers keep importing auth/provider concepts from `codex-core`. +- Only boundary crates may depend on `codex-auth` directly. + - Allowed: `codex-core`, `codex-login`, `codex-auth` + - Not allowed: `tui`, `exec`, `app-server`, `chatgpt`, `cli` +- Avoid duplicating logic between `codex-core` and `codex-auth`. +- Prefer moving code first and redesigning APIs later. + +## Ownership Boundary + +### `codex-auth` owns + +- auth state and persistence +- token parsing and refresh +- auth manager lifecycle +- auth storage backends +- model provider definitions +- built-in provider registry +- auth/provider request glue +- provider-aware remote model catalog behavior if it is truly auth/provider infrastructure + +### `codex-core` owns + +- config loading and orchestration +- app/session/thread behavior +- product policy +- model presentation and presets +- thin re-exports over `codex-auth` + +## PR 1 + +### Scope + +Create `codex-auth`, move foundational auth/provider code into it, and add `codex-core` re-exports in the same PR. + +### Expected moves + +- `core/src/token_data.rs` +- `core/src/model_provider_info.rs` + +These should become `codex-auth` modules with minimal behavioral change. + +Note: +`auth/storage` and the heavier auth manager implementation may stay in `codex-core` for PR 1 if moving them would drag `Config` or other core-only seams into `codex-auth` too early. + +### Required facade work in the same PR + +- Re-export moved items from `codex-core` +- Preserve existing import paths where practical +- Keep module-path compatibility for callers that currently use: + - `codex_core::auth::...` + - `codex_core::model_provider_info::...` + +### Acceptance criteria + +- `codex-auth` exists and owns the moved implementation +- `codex-core` still exposes the same public auth/provider concepts +- No normal consumer crate imports `codex-auth` +- No behavior change intended + +### Checkpoints + +- [ ] `core/src/lib.rs` re-exports moved auth/provider symbols +- [ ] compatibility module exists for `codex_core::model_provider_info::*` +- [ ] `core/Cargo.toml` depends on `codex-auth` +- [ ] no direct `codex-auth` dependency added to `tui`, `exec`, `app-server`, `chatgpt`, or `cli` +- [ ] moved code does not pull `codex-core` back into `codex-auth` + +## PR 2 + +### Scope + +Move implementation-boundary consumers and auth/provider glue to the new owner while keeping `codex-core` as the facade. + +### Expected changes + +- move `codex-login` to depend on `codex-auth` directly +- move auth/provider request glue that belongs with auth ownership +- update `codex-core` internals to consume `codex-auth` directly under the facade + +### Likely files to review + +- `login/src/lib.rs` +- `login/Cargo.toml` +- `core/src/api_bridge.rs` +- `core/src/client.rs` +- `core/src/thread_manager.rs` + +### Acceptance criteria + +- `codex-login` no longer depends on `codex-core` for auth ownership +- auth/provider glue lives with `codex-auth` +- public consumption for normal crates still goes through `codex-core` + +### Checkpoints + +- [ ] `login` depends on `codex-auth`, not `codex-core`, for auth concepts +- [ ] `core` still re-exports the same public surface +- [ ] no normal consumer crate was switched to importing `codex-auth` +- [ ] auth header/account-id/provider glue no longer has split ownership +- [ ] dependency direction is improved, not inverted again + +## PR 3 + +### Scope + +Move the auth/provider-owned model-management pieces and finish cleanup without dragging product policy out of `codex-core`. + +### Expected changes + +- move only the provider/auth-coupled parts of model management +- keep model presentation, presets, and app policy in `codex-core` +- remove leftover shims that are no longer needed, while keeping the public `codex-core` facade stable + +### Likely files to review + +- `core/src/models_manager/manager.rs` +- `core/src/models_manager/mod.rs` +- `core/src/models_manager/model_info.rs` + +### Acceptance criteria + +- provider/auth infrastructure no longer lives in `codex-core` +- `codex-core` still owns product policy and presentation +- normal consumer crates still do not need `codex-auth` + +### Checkpoints + +- [ ] only provider/auth infrastructure moved out of `models_manager` +- [ ] model presets and presentation remain in `codex-core` +- [ ] `codex-core` facade still covers auth/provider concepts for consumers +- [ ] no duplicate provider/auth logic remains across crates + +## Watchlist + +These are the highest-risk seams to review during the migration: + +- `core/src/api_bridge.rs` +- `core/src/client.rs` +- `core/src/thread_manager.rs` +- `core/src/config/mod.rs` +- `core/src/models_manager/manager.rs` +- `login/src/lib.rs` +- any `codex_core::auth::...` imports in non-`core` crates + +## Quick Rule For Future Work + +If code answers one of these questions, it likely belongs in `codex-auth`: + +- Who am I authenticated as? +- How are credentials stored or refreshed? +- What provider am I configured to use? +- How does auth interact with provider request construction? + +If code answers one of these questions, it likely stays in `codex-core`: + +- How should the app behave? +- How should models be presented to users? +- How should config be orchestrated for the product? +- How should threads, sessions, and UI flows be managed? diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d6a13a3d4c..188336720f 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1596,6 +1596,24 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "codex-auth" +version = "0.0.0" +dependencies = [ + "base64 0.22.1", + "chrono", + "codex-api", + "codex-app-server-protocol", + "http 1.4.0", + "maplit", + "pretty_assertions", + "schemars 0.8.22", + "serde", + "serde_json", + "thiserror 2.0.18", + "toml 0.9.11+spec-1.1.0", +] + [[package]] name = "codex-backend-client" version = "0.0.0" @@ -1840,6 +1858,7 @@ dependencies = [ "codex-arg0", "codex-artifacts", "codex-async-utils", + "codex-auth", "codex-client", "codex-config", "codex-connectors", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 35ff64195e..f705c41ceb 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -18,6 +18,7 @@ members = [ "cli", "connectors", "config", + "codex-auth", "shell-command", "shell-escalation", "skills", @@ -87,6 +88,7 @@ license = "Apache-2.0" app_test_support = { path = "app-server/tests/common" } codex-ansi-escape = { path = "ansi-escape" } codex-api = { path = "codex-api" } +codex-auth = { path = "codex-auth" } codex-artifacts = { path = "artifacts" } codex-package-manager = { path = "package-manager" } codex-app-server = { path = "app-server" } diff --git a/codex-rs/codex-auth/BUILD.bazel b/codex-rs/codex-auth/BUILD.bazel new file mode 100644 index 0000000000..293a40d209 --- /dev/null +++ b/codex-rs/codex-auth/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "codex-auth", + crate_name = "codex_auth", +) diff --git a/codex-rs/codex-auth/Cargo.toml b/codex-rs/codex-auth/Cargo.toml new file mode 100644 index 0000000000..be7265029f --- /dev/null +++ b/codex-rs/codex-auth/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "codex-auth" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lints] +workspace = true + +[dependencies] +base64 = { workspace = true } +chrono = { workspace = true, features = ["serde"] } +codex-api = { workspace = true } +codex-app-server-protocol = { workspace = true } +http = { workspace = true } +schemars = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +thiserror = { workspace = true } + +[dev-dependencies] +maplit = { workspace = true } +pretty_assertions = { workspace = true } +toml = { workspace = true } diff --git a/codex-rs/codex-auth/src/error.rs b/codex-rs/codex-auth/src/error.rs new file mode 100644 index 0000000000..94edf8765a --- /dev/null +++ b/codex-rs/codex-auth/src/error.rs @@ -0,0 +1,15 @@ +#[derive(Debug)] +pub struct EnvVarError { + pub var: String, + pub instructions: Option, +} + +impl std::fmt::Display for EnvVarError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "Missing environment variable: `{}`.", self.var)?; + if let Some(instructions) = &self.instructions { + write!(f, " {instructions}")?; + } + Ok(()) + } +} diff --git a/codex-rs/codex-auth/src/lib.rs b/codex-rs/codex-auth/src/lib.rs new file mode 100644 index 0000000000..78f9bfa132 --- /dev/null +++ b/codex-rs/codex-auth/src/lib.rs @@ -0,0 +1,19 @@ +pub mod error; +pub mod provider; +pub mod token_data; + +#[cfg(test)] +mod model_provider_info_tests; +#[cfg(test)] +mod token_data_tests; + +pub use error::EnvVarError; +pub use provider::DEFAULT_LMSTUDIO_PORT; +pub use provider::DEFAULT_OLLAMA_PORT; +pub use provider::LMSTUDIO_OSS_PROVIDER_ID; +pub use provider::ModelProviderInfo; +pub use provider::OLLAMA_OSS_PROVIDER_ID; +pub use provider::OPENAI_PROVIDER_ID; +pub use provider::WireApi; +pub use provider::built_in_model_providers; +pub use provider::create_oss_provider_with_base_url; diff --git a/codex-rs/core/src/model_provider_info_tests.rs b/codex-rs/codex-auth/src/model_provider_info_tests.rs similarity index 99% rename from codex-rs/core/src/model_provider_info_tests.rs rename to codex-rs/codex-auth/src/model_provider_info_tests.rs index a5309117ae..d62bf7f193 100644 --- a/codex-rs/core/src/model_provider_info_tests.rs +++ b/codex-rs/codex-auth/src/model_provider_info_tests.rs @@ -1,4 +1,4 @@ -use super::*; +use super::provider::*; use pretty_assertions::assert_eq; #[test] diff --git a/codex-rs/codex-auth/src/provider.rs b/codex-rs/codex-auth/src/provider.rs new file mode 100644 index 0000000000..c1b25529bd --- /dev/null +++ b/codex-rs/codex-auth/src/provider.rs @@ -0,0 +1,286 @@ +use crate::error::EnvVarError; +use codex_api::Provider as ApiProvider; +use codex_api::provider::RetryConfig as ApiRetryConfig; +use codex_app_server_protocol::AuthMode as ApiAuthMode; +use http::HeaderMap; +use http::header::HeaderName; +use http::header::HeaderValue; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; +use std::collections::HashMap; +use std::fmt; +use std::time::Duration; + +const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; +const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; +const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; +pub const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS: u64 = 15_000; +const MAX_STREAM_MAX_RETRIES: u64 = 100; +const MAX_REQUEST_MAX_RETRIES: u64 = 100; + +const OPENAI_PROVIDER_NAME: &str = "OpenAI"; +pub const OPENAI_PROVIDER_ID: &str = "openai"; +pub const CHAT_WIRE_API_REMOVED_ERROR: &str = "`wire_api = \"chat\"` is no longer supported.\nHow to fix: set `wire_api = \"responses\"` in your provider config.\nMore info: https://github.com/openai/codex/discussions/7782"; +pub const LEGACY_OLLAMA_CHAT_PROVIDER_ID: &str = "ollama-chat"; +pub const OLLAMA_CHAT_PROVIDER_REMOVED_ERROR: &str = "`ollama-chat` is no longer supported.\nHow to fix: replace `ollama-chat` with `ollama` in `model_provider`, `oss_provider`, or `--local-provider`.\nMore info: https://github.com/openai/codex/discussions/7782"; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, JsonSchema)] +#[serde(rename_all = "lowercase")] +pub enum WireApi { + #[default] + Responses, +} + +impl fmt::Display for WireApi { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let value = match self { + Self::Responses => "responses", + }; + f.write_str(value) + } +} + +impl<'de> Deserialize<'de> for WireApi { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + match value.as_str() { + "responses" => Ok(Self::Responses), + "chat" => Err(serde::de::Error::custom(CHAT_WIRE_API_REMOVED_ERROR)), + _ => Err(serde::de::Error::unknown_variant(&value, &["responses"])), + } + } +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct ModelProviderInfo { + pub name: String, + pub base_url: Option, + pub env_key: Option, + pub env_key_instructions: Option, + pub experimental_bearer_token: Option, + #[serde(default)] + pub wire_api: WireApi, + pub query_params: Option>, + pub http_headers: Option>, + pub env_http_headers: Option>, + pub request_max_retries: Option, + pub stream_max_retries: Option, + pub stream_idle_timeout_ms: Option, + pub websocket_connect_timeout_ms: Option, + #[serde(default)] + pub requires_openai_auth: bool, + #[serde(default)] + pub supports_websockets: bool, +} + +impl ModelProviderInfo { + fn build_header_map(&self) -> HeaderMap { + let capacity = self.http_headers.as_ref().map_or(0, HashMap::len) + + self.env_http_headers.as_ref().map_or(0, HashMap::len); + let mut headers = HeaderMap::with_capacity(capacity); + if let Some(extra) = &self.http_headers { + for (k, v) in extra { + if let (Ok(name), Ok(value)) = (HeaderName::try_from(k), HeaderValue::try_from(v)) { + headers.insert(name, value); + } + } + } + + if let Some(env_headers) = &self.env_http_headers { + for (header, env_var) in env_headers { + if let Ok(val) = std::env::var(env_var) + && !val.trim().is_empty() + && let (Ok(name), Ok(value)) = + (HeaderName::try_from(header), HeaderValue::try_from(val)) + { + headers.insert(name, value); + } + } + } + + headers + } + + pub fn to_api_provider( + &self, + auth_mode: Option, + ) -> Result { + let default_base_url = if matches!( + auth_mode, + Some(ApiAuthMode::Chatgpt | ApiAuthMode::ChatgptAuthTokens) + ) { + "https://chatgpt.com/backend-api/codex" + } else { + "https://api.openai.com/v1" + }; + let base_url = self + .base_url + .clone() + .unwrap_or_else(|| default_base_url.to_string()); + + let retry = ApiRetryConfig { + max_attempts: self.request_max_retries(), + base_delay: Duration::from_millis(200), + retry_429: false, + retry_5xx: true, + retry_transport: true, + }; + + Ok(ApiProvider { + name: self.name.clone(), + base_url, + query_params: self.query_params.clone(), + headers: self.build_header_map(), + retry, + stream_idle_timeout: self.stream_idle_timeout(), + }) + } + + pub fn api_key(&self) -> Result, EnvVarError> { + match &self.env_key { + Some(env_key) => { + let api_key = std::env::var(env_key) + .ok() + .filter(|v| !v.trim().is_empty()) + .ok_or_else(|| EnvVarError { + var: env_key.clone(), + instructions: self.env_key_instructions.clone(), + })?; + Ok(Some(api_key)) + } + None => Ok(None), + } + } + + pub fn request_max_retries(&self) -> u64 { + self.request_max_retries + .unwrap_or(DEFAULT_REQUEST_MAX_RETRIES) + .min(MAX_REQUEST_MAX_RETRIES) + } + + pub fn stream_max_retries(&self) -> u64 { + self.stream_max_retries + .unwrap_or(DEFAULT_STREAM_MAX_RETRIES) + .min(MAX_STREAM_MAX_RETRIES) + } + + pub fn stream_idle_timeout(&self) -> Duration { + self.stream_idle_timeout_ms + .map(Duration::from_millis) + .unwrap_or(Duration::from_millis(DEFAULT_STREAM_IDLE_TIMEOUT_MS)) + } + + pub fn websocket_connect_timeout(&self) -> Duration { + self.websocket_connect_timeout_ms + .map(Duration::from_millis) + .unwrap_or(Duration::from_millis(DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS)) + } + + pub fn create_openai_provider(base_url: Option) -> ModelProviderInfo { + ModelProviderInfo { + name: OPENAI_PROVIDER_NAME.into(), + base_url, + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + wire_api: WireApi::Responses, + query_params: None, + http_headers: Some( + [("version".to_string(), env!("CARGO_PKG_VERSION").to_string())] + .into_iter() + .collect(), + ), + env_http_headers: Some( + [ + ( + "OpenAI-Organization".to_string(), + "OPENAI_ORGANIZATION".to_string(), + ), + ("OpenAI-Project".to_string(), "OPENAI_PROJECT".to_string()), + ] + .into_iter() + .collect(), + ), + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: true, + supports_websockets: true, + } + } + + pub fn is_openai(&self) -> bool { + self.name == OPENAI_PROVIDER_NAME + } +} + +pub const DEFAULT_LMSTUDIO_PORT: u16 = 1234; +pub const DEFAULT_OLLAMA_PORT: u16 = 11434; + +pub const LMSTUDIO_OSS_PROVIDER_ID: &str = "lmstudio"; +pub const OLLAMA_OSS_PROVIDER_ID: &str = "ollama"; + +pub fn built_in_model_providers( + openai_base_url: Option, +) -> HashMap { + use ModelProviderInfo as P; + let openai_provider = P::create_openai_provider(openai_base_url); + + [ + (OPENAI_PROVIDER_ID, openai_provider), + ( + OLLAMA_OSS_PROVIDER_ID, + create_oss_provider(DEFAULT_OLLAMA_PORT, WireApi::Responses), + ), + ( + LMSTUDIO_OSS_PROVIDER_ID, + create_oss_provider(DEFAULT_LMSTUDIO_PORT, WireApi::Responses), + ), + ] + .into_iter() + .map(|(k, v)| (k.to_string(), v)) + .collect() +} + +pub fn create_oss_provider(default_provider_port: u16, wire_api: WireApi) -> ModelProviderInfo { + let default_codex_oss_base_url = format!( + "http://localhost:{codex_oss_port}/v1", + codex_oss_port = std::env::var("CODEX_OSS_PORT") + .ok() + .filter(|value| !value.trim().is_empty()) + .and_then(|value| value.parse::().ok()) + .unwrap_or(default_provider_port) + ); + + let codex_oss_base_url = std::env::var("CODEX_OSS_BASE_URL") + .ok() + .filter(|v| !v.trim().is_empty()) + .unwrap_or(default_codex_oss_base_url); + create_oss_provider_with_base_url(&codex_oss_base_url, wire_api) +} + +pub fn create_oss_provider_with_base_url(base_url: &str, wire_api: WireApi) -> ModelProviderInfo { + ModelProviderInfo { + name: "gpt-oss".into(), + base_url: Some(base_url.into()), + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + wire_api, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + } +} diff --git a/codex-rs/codex-auth/src/token_data.rs b/codex-rs/codex-auth/src/token_data.rs new file mode 100644 index 0000000000..392b4c22c7 --- /dev/null +++ b/codex-rs/codex-auth/src/token_data.rs @@ -0,0 +1,163 @@ +use base64::Engine; +use serde::Deserialize; +use serde::Serialize; +use thiserror::Error; + +#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Default)] +pub struct TokenData { + #[serde( + deserialize_with = "deserialize_id_token", + serialize_with = "serialize_id_token" + )] + pub id_token: IdTokenInfo, + pub access_token: String, + pub refresh_token: String, + pub account_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] +pub struct IdTokenInfo { + pub email: Option, + pub chatgpt_plan_type: Option, + pub chatgpt_user_id: Option, + pub chatgpt_account_id: Option, + pub raw_jwt: String, +} + +impl IdTokenInfo { + pub fn get_chatgpt_plan_type(&self) -> Option { + self.chatgpt_plan_type.as_ref().map(|t| match t { + PlanType::Known(plan) => format!("{plan:?}"), + PlanType::Unknown(s) => s.clone(), + }) + } + + pub fn is_workspace_account(&self) -> bool { + matches!( + self.chatgpt_plan_type, + Some(PlanType::Known( + KnownPlan::Team | KnownPlan::Business | KnownPlan::Enterprise | KnownPlan::Edu + )) + ) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum PlanType { + Known(KnownPlan), + Unknown(String), +} + +impl PlanType { + pub fn from_raw_value(raw: &str) -> Self { + match raw.to_ascii_lowercase().as_str() { + "free" => Self::Known(KnownPlan::Free), + "go" => Self::Known(KnownPlan::Go), + "plus" => Self::Known(KnownPlan::Plus), + "pro" => Self::Known(KnownPlan::Pro), + "team" => Self::Known(KnownPlan::Team), + "business" => Self::Known(KnownPlan::Business), + "enterprise" => Self::Known(KnownPlan::Enterprise), + "education" | "edu" => Self::Known(KnownPlan::Edu), + _ => Self::Unknown(raw.to_string()), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum KnownPlan { + Free, + Go, + Plus, + Pro, + Team, + Business, + Enterprise, + Edu, +} + +#[derive(Deserialize)] +struct IdClaims { + #[serde(default)] + email: Option, + #[serde(rename = "https://api.openai.com/profile", default)] + profile: Option, + #[serde(rename = "https://api.openai.com/auth", default)] + auth: Option, +} + +#[derive(Deserialize)] +struct ProfileClaims { + #[serde(default)] + email: Option, +} + +#[derive(Deserialize)] +struct AuthClaims { + #[serde(default)] + chatgpt_plan_type: Option, + #[serde(default)] + chatgpt_user_id: Option, + #[serde(default)] + user_id: Option, + #[serde(default)] + chatgpt_account_id: Option, +} + +#[derive(Debug, Error)] +pub enum IdTokenInfoError { + #[error("invalid ID token format")] + InvalidFormat, + #[error(transparent)] + Base64(#[from] base64::DecodeError), + #[error(transparent)] + Json(#[from] serde_json::Error), +} + +pub fn parse_chatgpt_jwt_claims(jwt: &str) -> Result { + let mut parts = jwt.split('.'); + let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) { + (Some(h), Some(p), Some(s)) if !h.is_empty() && !p.is_empty() && !s.is_empty() => (h, p, s), + _ => return Err(IdTokenInfoError::InvalidFormat), + }; + + let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(payload_b64)?; + let claims: IdClaims = serde_json::from_slice(&payload_bytes)?; + let email = claims + .email + .or_else(|| claims.profile.and_then(|profile| profile.email)); + + match claims.auth { + Some(auth) => Ok(IdTokenInfo { + email, + raw_jwt: jwt.to_string(), + chatgpt_plan_type: auth.chatgpt_plan_type, + chatgpt_user_id: auth.chatgpt_user_id.or(auth.user_id), + chatgpt_account_id: auth.chatgpt_account_id, + }), + None => Ok(IdTokenInfo { + email, + raw_jwt: jwt.to_string(), + chatgpt_plan_type: None, + chatgpt_user_id: None, + chatgpt_account_id: None, + }), + } +} + +fn deserialize_id_token<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + let s = String::deserialize(deserializer)?; + parse_chatgpt_jwt_claims(&s).map_err(serde::de::Error::custom) +} + +fn serialize_id_token(id_token: &IdTokenInfo, serializer: S) -> Result +where + S: serde::Serializer, +{ + serializer.serialize_str(&id_token.raw_jwt) +} diff --git a/codex-rs/core/src/token_data_tests.rs b/codex-rs/codex-auth/src/token_data_tests.rs similarity index 98% rename from codex-rs/core/src/token_data_tests.rs rename to codex-rs/codex-auth/src/token_data_tests.rs index e599379c18..df8f7353d0 100644 --- a/codex-rs/core/src/token_data_tests.rs +++ b/codex-rs/codex-auth/src/token_data_tests.rs @@ -1,4 +1,5 @@ -use super::*; +use base64::Engine; +use super::token_data::*; use pretty_assertions::assert_eq; use serde::Serialize; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index d11e209813..7a2f6d61ba 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -28,6 +28,7 @@ chardetng = { workspace = true } chrono = { workspace = true, features = ["serde"] } clap = { workspace = true, features = ["derive"] } codex-api = { workspace = true } +codex-auth = { workspace = true } codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } codex-async-utils = { workspace = true } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index eb5cb4c08a..3b1f501974 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -528,7 +528,7 @@ impl ModelClient { let api_provider = self .state .provider - .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?; + .to_api_provider(auth.as_ref().map(CodexAuth::api_auth_mode))?; let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?; Ok(CurrentClientSetup { auth, diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index e8e86defc2..6d218c54f9 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -9,6 +9,7 @@ use chrono::Datelike; use chrono::Local; use chrono::Utc; use codex_async_utils::CancelErr; +use codex_auth::EnvVarError; use codex_protocol::ThreadId; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::ErrorEvent; @@ -191,6 +192,12 @@ impl From for CodexErr { } } +impl From for CodexErr { + fn from(error: EnvVarError) -> Self { + Self::EnvVar(error) + } +} + impl CodexErr { pub fn is_retryable(&self) -> bool { match self { @@ -551,26 +558,6 @@ fn now_for_retry() -> DateTime { Utc::now() } -#[derive(Debug)] -pub struct EnvVarError { - /// Name of the environment variable that is missing. - pub var: String, - - /// Optional instructions to help the user get a valid value for the - /// variable and set it. - pub instructions: Option, -} - -impl std::fmt::Display for EnvVarError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "Missing environment variable: `{}`.", self.var)?; - if let Some(instructions) = &self.instructions { - write!(f, " {instructions}")?; - } - Ok(()) - } -} - impl CodexErr { /// Minimal shim so that existing `e.downcast_ref::()` checks continue to compile /// after replacing `anyhow::Error` in the return signature. This mirrors the behavior of diff --git a/codex-rs/core/src/model_provider_info.rs b/codex-rs/core/src/model_provider_info.rs index 737a47780d..a99f37c209 100644 --- a/codex-rs/core/src/model_provider_info.rs +++ b/codex-rs/core/src/model_provider_info.rs @@ -1,356 +1 @@ -//! Registry of model providers supported by Codex. -//! -//! Providers can be defined in two places: -//! 1. Built-in defaults compiled into the binary so Codex works out-of-the-box. -//! 2. User-defined entries inside `~/.codex/config.toml` under the `model_providers` -//! key. These override or extend the defaults at runtime. - -use crate::auth::AuthMode; -use crate::error::EnvVarError; -use codex_api::Provider as ApiProvider; -use codex_api::provider::RetryConfig as ApiRetryConfig; -use http::HeaderMap; -use http::header::HeaderName; -use http::header::HeaderValue; -use schemars::JsonSchema; -use serde::Deserialize; -use serde::Serialize; -use std::collections::HashMap; -use std::fmt; -use std::time::Duration; - -const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; -const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; -const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; -pub(crate) const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS: u64 = 15_000; -/// Hard cap for user-configured `stream_max_retries`. -const MAX_STREAM_MAX_RETRIES: u64 = 100; -/// Hard cap for user-configured `request_max_retries`. -const MAX_REQUEST_MAX_RETRIES: u64 = 100; - -const OPENAI_PROVIDER_NAME: &str = "OpenAI"; -pub const OPENAI_PROVIDER_ID: &str = "openai"; -const CHAT_WIRE_API_REMOVED_ERROR: &str = "`wire_api = \"chat\"` is no longer supported.\nHow to fix: set `wire_api = \"responses\"` in your provider config.\nMore info: https://github.com/openai/codex/discussions/7782"; -pub(crate) const LEGACY_OLLAMA_CHAT_PROVIDER_ID: &str = "ollama-chat"; -pub(crate) const OLLAMA_CHAT_PROVIDER_REMOVED_ERROR: &str = "`ollama-chat` is no longer supported.\nHow to fix: replace `ollama-chat` with `ollama` in `model_provider`, `oss_provider`, or `--local-provider`.\nMore info: https://github.com/openai/codex/discussions/7782"; - -/// Wire protocol that the provider speaks. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, JsonSchema)] -#[serde(rename_all = "lowercase")] -pub enum WireApi { - /// The Responses API exposed by OpenAI at `/v1/responses`. - #[default] - Responses, -} - -impl fmt::Display for WireApi { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let value = match self { - Self::Responses => "responses", - }; - f.write_str(value) - } -} - -impl<'de> Deserialize<'de> for WireApi { - fn deserialize(deserializer: D) -> Result - where - D: serde::Deserializer<'de>, - { - let value = String::deserialize(deserializer)?; - match value.as_str() { - "responses" => Ok(Self::Responses), - "chat" => Err(serde::de::Error::custom(CHAT_WIRE_API_REMOVED_ERROR)), - _ => Err(serde::de::Error::unknown_variant(&value, &["responses"])), - } - } -} - -/// Serializable representation of a provider definition. -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)] -#[schemars(deny_unknown_fields)] -pub struct ModelProviderInfo { - /// Friendly display name. - pub name: String, - /// Base URL for the provider's OpenAI-compatible API. - pub base_url: Option, - /// Environment variable that stores the user's API key for this provider. - pub env_key: Option, - - /// Optional instructions to help the user get a valid value for the - /// variable and set it. - pub env_key_instructions: Option, - - /// Value to use with `Authorization: Bearer ` header. Use of this - /// config is discouraged in favor of `env_key` for security reasons, but - /// this may be necessary when using this programmatically. - pub experimental_bearer_token: Option, - - /// Which wire protocol this provider expects. - #[serde(default)] - pub wire_api: WireApi, - - /// Optional query parameters to append to the base URL. - pub query_params: Option>, - - /// Additional HTTP headers to include in requests to this provider where - /// the (key, value) pairs are the header name and value. - pub http_headers: Option>, - - /// Optional HTTP headers to include in requests to this provider where the - /// (key, value) pairs are the header name and _environment variable_ whose - /// value should be used. If the environment variable is not set, or the - /// value is empty, the header will not be included in the request. - pub env_http_headers: Option>, - - /// Maximum number of times to retry a failed HTTP request to this provider. - pub request_max_retries: Option, - - /// Number of times to retry reconnecting a dropped streaming response before failing. - pub stream_max_retries: Option, - - /// Idle timeout (in milliseconds) to wait for activity on a streaming response before treating - /// the connection as lost. - pub stream_idle_timeout_ms: Option, - - /// Maximum time (in milliseconds) to wait for a websocket connection attempt before treating - /// it as failed. - pub websocket_connect_timeout_ms: Option, - - /// Does this provider require an OpenAI API Key or ChatGPT login token? If true, - /// user is presented with login screen on first run, and login preference and token/key - /// are stored in auth.json. If false (which is the default), login screen is skipped, - /// and API key (if needed) comes from the "env_key" environment variable. - #[serde(default)] - pub requires_openai_auth: bool, - - /// Whether this provider supports the Responses API WebSocket transport. - #[serde(default)] - pub supports_websockets: bool, -} - -impl ModelProviderInfo { - fn build_header_map(&self) -> crate::error::Result { - let capacity = self.http_headers.as_ref().map_or(0, HashMap::len) - + self.env_http_headers.as_ref().map_or(0, HashMap::len); - let mut headers = HeaderMap::with_capacity(capacity); - if let Some(extra) = &self.http_headers { - for (k, v) in extra { - if let (Ok(name), Ok(value)) = (HeaderName::try_from(k), HeaderValue::try_from(v)) { - headers.insert(name, value); - } - } - } - - if let Some(env_headers) = &self.env_http_headers { - for (header, env_var) in env_headers { - if let Ok(val) = std::env::var(env_var) - && !val.trim().is_empty() - && let (Ok(name), Ok(value)) = - (HeaderName::try_from(header), HeaderValue::try_from(val)) - { - headers.insert(name, value); - } - } - } - - Ok(headers) - } - - pub(crate) fn to_api_provider( - &self, - auth_mode: Option, - ) -> crate::error::Result { - let default_base_url = if matches!(auth_mode, Some(AuthMode::Chatgpt)) { - "https://chatgpt.com/backend-api/codex" - } else { - "https://api.openai.com/v1" - }; - let base_url = self - .base_url - .clone() - .unwrap_or_else(|| default_base_url.to_string()); - - let headers = self.build_header_map()?; - let retry = ApiRetryConfig { - max_attempts: self.request_max_retries(), - base_delay: Duration::from_millis(200), - retry_429: false, - retry_5xx: true, - retry_transport: true, - }; - - Ok(ApiProvider { - name: self.name.clone(), - base_url, - query_params: self.query_params.clone(), - headers, - retry, - stream_idle_timeout: self.stream_idle_timeout(), - }) - } - - /// If `env_key` is Some, returns the API key for this provider if present - /// (and non-empty) in the environment. If `env_key` is required but - /// cannot be found, returns an error. - pub fn api_key(&self) -> crate::error::Result> { - match &self.env_key { - Some(env_key) => { - let api_key = std::env::var(env_key) - .ok() - .filter(|v| !v.trim().is_empty()) - .ok_or_else(|| { - crate::error::CodexErr::EnvVar(EnvVarError { - var: env_key.clone(), - instructions: self.env_key_instructions.clone(), - }) - })?; - Ok(Some(api_key)) - } - None => Ok(None), - } - } - - /// Effective maximum number of request retries for this provider. - pub fn request_max_retries(&self) -> u64 { - self.request_max_retries - .unwrap_or(DEFAULT_REQUEST_MAX_RETRIES) - .min(MAX_REQUEST_MAX_RETRIES) - } - - /// Effective maximum number of stream reconnection attempts for this provider. - pub fn stream_max_retries(&self) -> u64 { - self.stream_max_retries - .unwrap_or(DEFAULT_STREAM_MAX_RETRIES) - .min(MAX_STREAM_MAX_RETRIES) - } - - /// Effective idle timeout for streaming responses. - pub fn stream_idle_timeout(&self) -> Duration { - self.stream_idle_timeout_ms - .map(Duration::from_millis) - .unwrap_or(Duration::from_millis(DEFAULT_STREAM_IDLE_TIMEOUT_MS)) - } - - /// Effective timeout for websocket connect attempts. - pub fn websocket_connect_timeout(&self) -> Duration { - self.websocket_connect_timeout_ms - .map(Duration::from_millis) - .unwrap_or(Duration::from_millis(DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS)) - } - - pub fn create_openai_provider(base_url: Option) -> ModelProviderInfo { - ModelProviderInfo { - name: OPENAI_PROVIDER_NAME.into(), - base_url, - env_key: None, - env_key_instructions: None, - experimental_bearer_token: None, - wire_api: WireApi::Responses, - query_params: None, - http_headers: Some( - [("version".to_string(), env!("CARGO_PKG_VERSION").to_string())] - .into_iter() - .collect(), - ), - env_http_headers: Some( - [ - ( - "OpenAI-Organization".to_string(), - "OPENAI_ORGANIZATION".to_string(), - ), - ("OpenAI-Project".to_string(), "OPENAI_PROJECT".to_string()), - ] - .into_iter() - .collect(), - ), - // Use global defaults for retry/timeout unless overridden in config.toml. - request_max_retries: None, - stream_max_retries: None, - stream_idle_timeout_ms: None, - websocket_connect_timeout_ms: None, - requires_openai_auth: true, - supports_websockets: true, - } - } - - pub fn is_openai(&self) -> bool { - self.name == OPENAI_PROVIDER_NAME - } -} - -pub const DEFAULT_LMSTUDIO_PORT: u16 = 1234; -pub const DEFAULT_OLLAMA_PORT: u16 = 11434; - -pub const LMSTUDIO_OSS_PROVIDER_ID: &str = "lmstudio"; -pub const OLLAMA_OSS_PROVIDER_ID: &str = "ollama"; - -/// Built-in default provider list. -pub fn built_in_model_providers( - openai_base_url: Option, -) -> HashMap { - use ModelProviderInfo as P; - let openai_provider = P::create_openai_provider(openai_base_url); - - // We do not want to be in the business of adjucating which third-party - // providers are bundled with Codex CLI, so we only include the OpenAI and - // open source ("oss") providers by default. Users are encouraged to add to - // `model_providers` in config.toml to add their own providers. - [ - (OPENAI_PROVIDER_ID, openai_provider), - ( - OLLAMA_OSS_PROVIDER_ID, - create_oss_provider(DEFAULT_OLLAMA_PORT, WireApi::Responses), - ), - ( - LMSTUDIO_OSS_PROVIDER_ID, - create_oss_provider(DEFAULT_LMSTUDIO_PORT, WireApi::Responses), - ), - ] - .into_iter() - .map(|(k, v)| (k.to_string(), v)) - .collect() -} - -pub fn create_oss_provider(default_provider_port: u16, wire_api: WireApi) -> ModelProviderInfo { - // These CODEX_OSS_ environment variables are experimental: we may - // switch to reading values from config.toml instead. - let default_codex_oss_base_url = format!( - "http://localhost:{codex_oss_port}/v1", - codex_oss_port = std::env::var("CODEX_OSS_PORT") - .ok() - .filter(|value| !value.trim().is_empty()) - .and_then(|value| value.parse::().ok()) - .unwrap_or(default_provider_port) - ); - - let codex_oss_base_url = std::env::var("CODEX_OSS_BASE_URL") - .ok() - .filter(|v| !v.trim().is_empty()) - .unwrap_or(default_codex_oss_base_url); - create_oss_provider_with_base_url(&codex_oss_base_url, wire_api) -} - -pub fn create_oss_provider_with_base_url(base_url: &str, wire_api: WireApi) -> ModelProviderInfo { - ModelProviderInfo { - name: "gpt-oss".into(), - base_url: Some(base_url.into()), - env_key: None, - env_key_instructions: None, - experimental_bearer_token: None, - wire_api, - query_params: None, - http_headers: None, - env_http_headers: None, - request_max_retries: None, - stream_max_retries: None, - stream_idle_timeout_ms: None, - websocket_connect_timeout_ms: None, - requires_openai_auth: false, - supports_websockets: false, - } -} - -#[cfg(test)] -#[path = "model_provider_info_tests.rs"] -mod tests; +pub use codex_auth::provider::*; diff --git a/codex-rs/core/src/models_manager/manager.rs b/codex-rs/core/src/models_manager/manager.rs index 29a1a85767..9d93b23487 100644 --- a/codex-rs/core/src/models_manager/manager.rs +++ b/codex-rs/core/src/models_manager/manager.rs @@ -433,7 +433,9 @@ impl ModelsManager { codex_otel::start_global_timer("codex.remote_models.fetch_update.duration_ms", &[]); let auth = self.auth_manager.auth().await; let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); - let api_provider = self.provider.to_api_provider(auth_mode)?; + let api_provider = self + .provider + .to_api_provider(auth.as_ref().map(CodexAuth::api_auth_mode))?; let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?; let auth_env = collect_auth_env_telemetry( &self.provider, diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index 1ddd72d0fd..15fd6ac878 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -22,6 +22,7 @@ use codex_api::RealtimeSessionMode; use codex_api::RealtimeWebsocketClient; use codex_api::endpoint::realtime_websocket::RealtimeWebsocketEvents; use codex_api::endpoint::realtime_websocket::RealtimeWebsocketWriter; +use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::ConversationAudioParams; use codex_protocol::protocol::ConversationStartParams; @@ -454,7 +455,7 @@ async fn prepare_realtime_start( let provider = sess.provider().await; let auth = sess.services.auth_manager.auth().await; let realtime_api_key = realtime_api_key(auth.as_ref(), &provider)?; - let mut api_provider = provider.to_api_provider(Some(crate::auth::AuthMode::ApiKey))?; + let mut api_provider = provider.to_api_provider(Some(ApiAuthMode::ApiKey))?; let config = sess.get_config().await; if let Some(realtime_ws_base_url) = &config.experimental_realtime_ws_base_url { api_provider.base_url = realtime_ws_base_url.clone(); diff --git a/codex-rs/core/src/token_data.rs b/codex-rs/core/src/token_data.rs index 5952d5940d..105e171a04 100644 --- a/codex-rs/core/src/token_data.rs +++ b/codex-rs/core/src/token_data.rs @@ -1,179 +1 @@ -use base64::Engine; -use serde::Deserialize; -use serde::Serialize; -use thiserror::Error; - -#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Default)] -pub struct TokenData { - /// Flat info parsed from the JWT in auth.json. - #[serde( - deserialize_with = "deserialize_id_token", - serialize_with = "serialize_id_token" - )] - pub id_token: IdTokenInfo, - - /// This is a JWT. - pub access_token: String, - - pub refresh_token: String, - - pub account_id: Option, -} - -/// Flat subset of useful claims in id_token from auth.json. -#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] -pub struct IdTokenInfo { - pub email: Option, - /// The ChatGPT subscription plan type - /// (e.g., "free", "plus", "pro", "business", "enterprise", "edu"). - /// (Note: values may vary by backend.) - pub(crate) chatgpt_plan_type: Option, - /// ChatGPT user identifier associated with the token, if present. - pub chatgpt_user_id: Option, - /// Organization/workspace identifier associated with the token, if present. - pub chatgpt_account_id: Option, - pub raw_jwt: String, -} - -impl IdTokenInfo { - pub fn get_chatgpt_plan_type(&self) -> Option { - self.chatgpt_plan_type.as_ref().map(|t| match t { - PlanType::Known(plan) => format!("{plan:?}"), - PlanType::Unknown(s) => s.clone(), - }) - } - - pub fn is_workspace_account(&self) -> bool { - matches!( - self.chatgpt_plan_type, - Some(PlanType::Known( - KnownPlan::Team | KnownPlan::Business | KnownPlan::Enterprise | KnownPlan::Edu - )) - ) - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(untagged)] -pub(crate) enum PlanType { - Known(KnownPlan), - Unknown(String), -} - -impl PlanType { - pub(crate) fn from_raw_value(raw: &str) -> Self { - match raw.to_ascii_lowercase().as_str() { - "free" => Self::Known(KnownPlan::Free), - "go" => Self::Known(KnownPlan::Go), - "plus" => Self::Known(KnownPlan::Plus), - "pro" => Self::Known(KnownPlan::Pro), - "team" => Self::Known(KnownPlan::Team), - "business" => Self::Known(KnownPlan::Business), - "enterprise" => Self::Known(KnownPlan::Enterprise), - "education" | "edu" => Self::Known(KnownPlan::Edu), - _ => Self::Unknown(raw.to_string()), - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub(crate) enum KnownPlan { - Free, - Go, - Plus, - Pro, - Team, - Business, - Enterprise, - Edu, -} - -#[derive(Deserialize)] -struct IdClaims { - #[serde(default)] - email: Option, - #[serde(rename = "https://api.openai.com/profile", default)] - profile: Option, - #[serde(rename = "https://api.openai.com/auth", default)] - auth: Option, -} - -#[derive(Deserialize)] -struct ProfileClaims { - #[serde(default)] - email: Option, -} - -#[derive(Deserialize)] -struct AuthClaims { - #[serde(default)] - chatgpt_plan_type: Option, - #[serde(default)] - chatgpt_user_id: Option, - #[serde(default)] - user_id: Option, - #[serde(default)] - chatgpt_account_id: Option, -} - -#[derive(Debug, Error)] -pub enum IdTokenInfoError { - #[error("invalid ID token format")] - InvalidFormat, - #[error(transparent)] - Base64(#[from] base64::DecodeError), - #[error(transparent)] - Json(#[from] serde_json::Error), -} - -pub fn parse_chatgpt_jwt_claims(jwt: &str) -> Result { - // JWT format: header.payload.signature - let mut parts = jwt.split('.'); - let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) { - (Some(h), Some(p), Some(s)) if !h.is_empty() && !p.is_empty() && !s.is_empty() => (h, p, s), - _ => return Err(IdTokenInfoError::InvalidFormat), - }; - - let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(payload_b64)?; - let claims: IdClaims = serde_json::from_slice(&payload_bytes)?; - let email = claims - .email - .or_else(|| claims.profile.and_then(|profile| profile.email)); - - match claims.auth { - Some(auth) => Ok(IdTokenInfo { - email, - raw_jwt: jwt.to_string(), - chatgpt_plan_type: auth.chatgpt_plan_type, - chatgpt_user_id: auth.chatgpt_user_id.or(auth.user_id), - chatgpt_account_id: auth.chatgpt_account_id, - }), - None => Ok(IdTokenInfo { - email, - raw_jwt: jwt.to_string(), - chatgpt_plan_type: None, - chatgpt_user_id: None, - chatgpt_account_id: None, - }), - } -} - -fn deserialize_id_token<'de, D>(deserializer: D) -> Result -where - D: serde::Deserializer<'de>, -{ - let s = String::deserialize(deserializer)?; - parse_chatgpt_jwt_claims(&s).map_err(serde::de::Error::custom) -} - -fn serialize_id_token(id_token: &IdTokenInfo, serializer: S) -> Result -where - S: serde::Serializer, -{ - serializer.serialize_str(&id_token.raw_jwt) -} - -#[cfg(test)] -#[path = "token_data_tests.rs"] -mod tests; +pub use codex_auth::token_data::*;