chore: proper client extraction (#6996)

2026-07-01 00:31:56 +08:00 · 2025-11-25 18:06:12 +00:00
parent 2845e2c006
commit 4502b1b263
45 changed files with 4890 additions and 2745 deletions
@@ -0,0 +1,154 @@
+use chrono::DateTime;
+use chrono::Utc;
+use codex_api::AuthProvider as ApiAuthProvider;
+use codex_api::TransportError;
+use codex_api::error::ApiError;
+use codex_api::rate_limits::parse_rate_limit;
+use http::HeaderMap;
+use serde::Deserialize;
+
+use crate::auth::CodexAuth;
+use crate::error::CodexErr;
+use crate::error::RetryLimitReachedError;
+use crate::error::UnexpectedResponseError;
+use crate::error::UsageLimitReachedError;
+use crate::model_provider_info::ModelProviderInfo;
+use crate::token_data::PlanType;
+
+pub(crate) fn map_api_error(err: ApiError) -> CodexErr {
+    match err {
+        ApiError::ContextWindowExceeded => CodexErr::ContextWindowExceeded,
+        ApiError::QuotaExceeded => CodexErr::QuotaExceeded,
+        ApiError::UsageNotIncluded => CodexErr::UsageNotIncluded,
+        ApiError::Retryable { message, delay } => CodexErr::Stream(message, delay),
+        ApiError::Stream(msg) => CodexErr::Stream(msg, None),
+        ApiError::Api { status, message } => CodexErr::UnexpectedStatus(UnexpectedResponseError {
+            status,
+            body: message,
+            request_id: None,
+        }),
+        ApiError::Transport(transport) => match transport {
+            TransportError::Http {
+                status,
+                headers,
+                body,
+            } => {
+                if status == http::StatusCode::INTERNAL_SERVER_ERROR {
+                    CodexErr::InternalServerError
+                } else if status == http::StatusCode::TOO_MANY_REQUESTS {
+                    if let Some(body) = body
+                        && let Ok(err) = serde_json::from_str::<UsageErrorResponse>(&body)
+                    {
+                        if err.error.error_type.as_deref() == Some("usage_limit_reached") {
+                            let rate_limits = headers.as_ref().and_then(parse_rate_limit);
+                            let resets_at = err
+                                .error
+                                .resets_at
+                                .and_then(|seconds| DateTime::<Utc>::from_timestamp(seconds, 0));
+                            return CodexErr::UsageLimitReached(UsageLimitReachedError {
+                                plan_type: err.error.plan_type,
+                                resets_at,
+                                rate_limits,
+                            });
+                        } else if err.error.error_type.as_deref() == Some("usage_not_included") {
+                            return CodexErr::UsageNotIncluded;
+                        }
+                    }
+
+                    CodexErr::RetryLimit(RetryLimitReachedError {
+                        status,
+                        request_id: extract_request_id(headers.as_ref()),
+                    })
+                } else {
+                    CodexErr::UnexpectedStatus(UnexpectedResponseError {
+                        status,
+                        body: body.unwrap_or_default(),
+                        request_id: extract_request_id(headers.as_ref()),
+                    })
+                }
+            }
+            TransportError::RetryLimit => CodexErr::RetryLimit(RetryLimitReachedError {
+                status: http::StatusCode::INTERNAL_SERVER_ERROR,
+                request_id: None,
+            }),
+            TransportError::Timeout => CodexErr::Timeout,
+            TransportError::Network(msg) | TransportError::Build(msg) => {
+                CodexErr::Stream(msg, None)
+            }
+        },
+        ApiError::RateLimit(msg) => CodexErr::Stream(msg, None),
+    }
+}
+
+fn extract_request_id(headers: Option<&HeaderMap>) -> Option<String> {
+    headers.and_then(|map| {
+        ["cf-ray", "x-request-id", "x-oai-request-id"]
+            .iter()
+            .find_map(|name| {
+                map.get(*name)
+                    .and_then(|v| v.to_str().ok())
+                    .map(str::to_string)
+            })
+    })
+}
+
+pub(crate) async fn auth_provider_from_auth(
+    auth: Option<CodexAuth>,
+    provider: &ModelProviderInfo,
+) -> crate::error::Result<CoreAuthProvider> {
+    if let Some(api_key) = provider.api_key()? {
+        return Ok(CoreAuthProvider {
+            token: Some(api_key),
+            account_id: None,
+        });
+    }
+
+    if let Some(token) = provider.experimental_bearer_token.clone() {
+        return Ok(CoreAuthProvider {
+            token: Some(token),
+            account_id: None,
+        });
+    }
+
+    if let Some(auth) = auth {
+        let token = auth.get_token().await?;
+        Ok(CoreAuthProvider {
+            token: Some(token),
+            account_id: auth.get_account_id(),
+        })
+    } else {
+        Ok(CoreAuthProvider {
+            token: None,
+            account_id: None,
+        })
+    }
+}
+
+#[derive(Debug, Deserialize)]
+struct UsageErrorResponse {
+    error: UsageErrorBody,
+}
+
+#[derive(Debug, Deserialize)]
+struct UsageErrorBody {
+    #[serde(rename = "type")]
+    error_type: Option<String>,
+    plan_type: Option<PlanType>,
+    resets_at: Option<i64>,
+}
+
+#[derive(Clone, Default)]
+pub(crate) struct CoreAuthProvider {
+    token: Option<String>,
+    account_id: Option<String>,
+}
+
+impl ApiAuthProvider for CoreAuthProvider {
+    fn bearer_token(&self) -> Option<String> {
+        self.token.clone()
+    }
+
+    fn account_id(&self) -> Option<String> {
+        self.account_id.clone()
+    }
+}