fix(core): align auth token schema and websocket errors

This commit is contained in:
James Peter
2026-02-10 19:28:57 +10:00
parent c5a6b184e0
commit 1261263044
5 changed files with 119 additions and 37 deletions
@@ -4,14 +4,20 @@
"accessToken": {
"type": "string"
},
"idToken": {
"chatgptAccountId": {
"type": "string"
},
"chatgptPlanType": {
"type": [
"string",
"null"
]
}
},
"required": [
"accessToken",
"idToken"
"chatgptAccountId"
],
"title": "ChatgptAuthTokensRefreshResponse",
"type": "object"
}
}
@@ -41,13 +41,20 @@
"description": "[UNSTABLE] FOR OPENAI INTERNAL USE ONLY - DO NOT USE. The access token must contain the same scopes that Codex-managed ChatGPT auth tokens have.",
"properties": {
"accessToken": {
"description": "Access token (JWT) supplied by the client. This token is used for backend API requests.",
"description": "Access token (JWT) supplied by the client. This token is used for backend API requests and email extraction.",
"type": "string"
},
"idToken": {
"description": "ID token (JWT) supplied by the client.\n\nThis token is used for identity and account metadata (email, plan type, workspace id).",
"chatgptAccountId": {
"description": "Workspace/account identifier supplied by the client.",
"type": "string"
},
"chatgptPlanType": {
"description": "Optional plan type supplied by the client.\n\nWhen `null`, Codex attempts to derive the plan type from access-token claims. If unavailable, the plan defaults to `unknown`.",
"type": [
"string",
"null"
]
},
"type": {
"enum": [
"chatgptAuthTokens"
@@ -58,7 +65,7 @@
},
"required": [
"accessToken",
"idToken",
"chatgptAccountId",
"type"
],
"title": "ChatgptAuthTokensv2::LoginAccountParams",
@@ -66,4 +73,4 @@
}
],
"title": "LoginAccountParams"
}
}
@@ -2,4 +2,4 @@
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
export type ChatgptAuthTokensRefreshResponse = { idToken: string, accessToken: string, };
export type ChatgptAuthTokensRefreshResponse = { accessToken: string, chatgptAccountId: string, chatgptPlanType: string | null, };
@@ -3,15 +3,19 @@
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
export type LoginAccountParams = { "type": "apiKey", apiKey: string, } | { "type": "chatgpt" } | { "type": "chatgptAuthTokens",
/**
* ID token (JWT) supplied by the client.
*
* This token is used for identity and account metadata (email, plan type,
* workspace id).
*/
idToken: string,
/**
* Access token (JWT) supplied by the client.
* This token is used for backend API requests.
* This token is used for backend API requests and email extraction.
*/
accessToken: string, };
accessToken: string,
/**
* Workspace/account identifier supplied by the client.
*/
chatgptAccountId: string,
/**
* Optional plan type supplied by the client.
*
* When `null`, Codex attempts to derive the plan type from access-token
* claims. If unavailable, the plan defaults to `unknown`.
*/
chatgptPlanType?: string | null, };
+84 -19
View File
@@ -33,6 +33,7 @@ use tracing::warn;
use uuid::Uuid;
use chrono::{DateTime, Duration as ChronoDuration, Utc};
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::Error as WsError;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
const AUTH_REQUIRED_MESSAGE: &str = "Authentication required. Run `code login` to continue.";
@@ -79,8 +80,6 @@ const RESPONSES_BETA_HEADER_V1: &str = "responses=v1";
const RESPONSES_BETA_HEADER_EXPERIMENTAL: &str = "responses=experimental";
const RESPONSES_WEBSOCKETS_BETA_HEADER: &str = "responses_websockets=2026-02-04";
const WEB_SEARCH_ELIGIBLE_HEADER: &str = "x-oai-web-search-eligible";
// Sticky-routing token captured at the start of a turn. When present, it must
// be replayed on every subsequent request within the same turn (retries,
// continuations, websocket reconnects).
@@ -102,6 +101,16 @@ struct ErrorResponse {
error: Error,
}
#[derive(Debug, Deserialize)]
struct WrappedWebsocketErrorEvent {
#[serde(rename = "type")]
kind: String,
#[serde(alias = "status_code")]
status: Option<u16>,
#[serde(default)]
error: Option<Error>,
}
#[derive(Debug, Deserialize)]
struct Error {
r#type: Option<String>,
@@ -657,7 +666,6 @@ impl ModelClient {
req_builder = attach_openai_subagent_header(req_builder);
req_builder = attach_codex_beta_features_header(req_builder, &self.config);
req_builder = attach_web_search_eligible_header(req_builder, &self.config);
if let Some(state) = turn_state.get() {
req_builder = req_builder.header(X_CODEX_TURN_STATE_HEADER, state);
}
@@ -803,6 +811,13 @@ impl ModelClient {
};
match next {
Ok(Message::Text(text)) => {
if let Some(error) = parse_wrapped_websocket_error_event(&text)
.and_then(map_wrapped_websocket_error_event)
{
let _ = tx_bytes.send(Err(error)).await;
break;
}
let chunk = format!("data: {text}\n\n");
if tx_bytes.send(Ok(Bytes::from(chunk))).await.is_err() {
break;
@@ -857,6 +872,11 @@ impl ModelClient {
return Ok(ResponseStream { rx_event });
}
Err(err) => {
if websocket_connect_is_upgrade_required(&err) {
warn!("responses websocket upgrade required; falling back to HTTP responses transport");
return self.stream_responses(prompt, log_tag).await;
}
let err = CodexErr::Stream(
format!("[ws] failed to connect: {err}"),
None,
@@ -1063,7 +1083,6 @@ impl ModelClient {
req_builder = attach_openai_subagent_header(req_builder);
req_builder = attach_codex_beta_features_header(req_builder, &self.config);
req_builder = attach_web_search_eligible_header(req_builder, &self.config);
if let Some(state) = turn_state.get() {
req_builder = req_builder.header(X_CODEX_TURN_STATE_HEADER, state);
}
@@ -1683,7 +1702,6 @@ impl ModelClient {
request = attach_openai_subagent_header(request);
request = attach_codex_beta_features_header(request, &self.config);
request = attach_web_search_eligible_header(request, &self.config);
if let Some(auth) = auth.as_ref()
&& auth.mode.is_chatgpt()
@@ -1822,24 +1840,71 @@ fn attach_codex_beta_features_header(
builder.header("x-codex-beta-features", value)
}
fn attach_web_search_eligible_header(
builder: reqwest::RequestBuilder,
config: &Config,
) -> reqwest::RequestBuilder {
let has_header = builder
.try_clone()
.and_then(|builder| builder.build().ok())
.map_or(false, |req| req.headers().contains_key(WEB_SEARCH_ELIGIBLE_HEADER));
if has_header {
return builder;
fn parse_wrapped_websocket_error_event(payload: &str) -> Option<WrappedWebsocketErrorEvent> {
let event: WrappedWebsocketErrorEvent = serde_json::from_str(payload).ok()?;
if event.kind != "error" {
return None;
}
Some(event)
}
fn map_wrapped_websocket_error_event(event: WrappedWebsocketErrorEvent) -> Option<CodexErr> {
let status = StatusCode::from_u16(event.status?).ok()?;
if status.is_success() {
return None;
}
let value = if config.tools_web_search_request {
"true"
let body = if let Some(error) = event.error {
if status == StatusCode::TOO_MANY_REQUESTS {
if error.r#type.as_deref() == Some("usage_limit_reached") {
return Some(CodexErr::UsageLimitReached(UsageLimitReachedError {
plan_type: error.plan_type,
resets_in_seconds: error.resets_in_seconds,
}));
}
if error.r#type.as_deref() == Some("usage_not_included") {
return Some(CodexErr::UsageNotIncluded);
}
}
if is_quota_exceeded_error(&error) {
return Some(CodexErr::QuotaExceeded);
}
serde_json::json!({
"error": {
"type": error.r#type,
"code": error.code,
"param": error.param,
"message": error.message,
"plan_type": error.plan_type,
"resets_in_seconds": error.resets_in_seconds,
}
})
.to_string()
} else {
"false"
serde_json::json!({
"error": {
"message": "websocket returned an error event"
}
})
.to_string()
};
builder.header(WEB_SEARCH_ELIGIBLE_HEADER, HeaderValue::from_static(value))
Some(CodexErr::UnexpectedStatus(UnexpectedResponseError {
status,
body,
request_id: None,
}))
}
fn websocket_connect_is_upgrade_required(error: &WsError) -> bool {
matches!(
error,
WsError::Http(response)
if response.status().as_u16() == 426
)
}
fn codex_beta_features_header_value(config: &Config) -> Option<HeaderValue> {