Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

第 5 章:API Client — 模型通信引擎

核心问题: OpenAI Codex CLI 如何与各种 LLM 提供商建立稳定、高效的通信?流式响应如何解析处理?认证系统如何设计?重试和限流策略如何实现?

5.1 架构概览:多提供商统一接口

OpenAI Codex CLI 的 API Client 系统设计为一个高度抽象的通信层,能够与多个 LLM 提供商无缝集成。其核心设计理念是 “Provider Agnostic”(提供商无关),通过统一接口屏蔽底层差异。

5.1.1 整体架构图

┌─────────────────────────────────────────────────────────────────┐
│                     ModelClient (Session Level)                │
├─────────────────────┬───────────────────────┬───────────────────┤
│   Authentication    │    Provider Router    │   Connection Pool │
│      Manager        │                       │                   │
├─────────────────────┼───────────────────────┼───────────────────┤
│                ModelClientSession (Turn Level)                 │
├─────────────────────┬───────────────────────┬───────────────────┤
│  Request Builder    │   Response Parser     │   Error Handler   │
└─────────────────────┴───────────────────────┴───────────────────┘
                                │
        ┌───────────────────────┼───────────────────────┐
        │                       │                       │
┌───────▼─────────┐   ┌─────────▼──────────┐   ┌───────▼─────────┐
│   OpenAI API    │   │   Responses API    │   │  Third-party    │
│   (Completions) │   │   (WebSocket)      │   │   Providers     │
└─────────────────┘   └────────────────────┘   └─────────────────┘

5.1.2 核心组件职责

在 client.rs 中定义的核心组件:

#![allow(unused)]
fn main() {
// codex-rs/core/src/client.rs

/// Session-scoped client for model provider APIs
/// 每个 Codex 会话持有一个 ModelClient 实例
pub struct ModelClient {
    auth_manager: Arc<AuthManager>,
    provider_config: ProviderConfig,
    connection_pool: ConnectionPool,
    request_telemetry: Arc<RequestTelemetry>,
    fallback_state: Arc<Mutex<FallbackState>>,
}

/// Turn-scoped session for streaming requests
/// 每个 Turn 创建一个 ModelClientSession
pub struct ModelClientSession {
    client: Arc<ModelClient>,
    turn_context: TurnContext,
    websocket_connection: Option<ApiWebSocketConnection>,
    turn_state_token: Option<String>, // 用于粘性路由
}

/// Provider configuration and capabilities
#[derive(Clone, Debug)]
pub struct ProviderConfig {
    provider_type: ProviderType,
    endpoint_url: String,
    supported_models: Vec<ModelInfo>,
    capabilities: ProviderCapabilities,
    rate_limits: RateLimitConfig,
}

#[derive(Debug, Clone)]
pub enum ProviderType {
    OpenAI,
    ResponsesApi, // OpenAI 内部 WebSocket API
    ThirdParty(String),
}
}

5.2 认证系统设计

5.2.1 多层认证架构

Codex 支持多种认证方式,以适应不同的部署环境和安全要求:

#![allow(unused)]
fn main() {
// codex-rs/core/src/auth.rs
#[derive(Debug, Clone)]
pub enum AuthMode {
    ApiKey(ApiKeyAuth),
    OAuth(OAuthAuth),
    ChatGptSession(SessionAuth),
    ServiceAccount(ServiceAccountAuth),
    DeviceCode(DeviceCodeAuth),
}

pub struct AuthManager {
    current_auth: Arc<RwLock<Option<CodexAuth>>>,
    auth_cache: Arc<Mutex<AuthCache>>,
    refresh_scheduler: Arc<RefreshScheduler>,
}

impl AuthManager {
    pub async fn get_valid_auth(&self) -> CodexResult<CodexAuth> {
        let current = self.current_auth.read().await;

        if let Some(auth) = current.as_ref() {
            if !auth.is_expired() {
                return Ok(auth.clone());
            }
        }

        drop(current);

        // Token 过期,尝试刷新
        self.refresh_auth().await
    }

    async fn refresh_auth(&self) -> CodexResult<CodexAuth> {
        let mut current = self.current_auth.write().await;

        // Double-check pattern,避免并发刷新
        if let Some(auth) = current.as_ref() {
            if !auth.is_expired() {
                return Ok(auth.clone());
            }
        }

        let refreshed = match current.as_ref() {
            Some(auth) => self.do_refresh_token(auth).await?,
            None => self.do_initial_auth().await?,
        };

        *current = Some(refreshed.clone());
        Ok(refreshed)
    }
}
}

5.2.2 OAuth 流程实现

对于 ChatGPT 等需要 OAuth 认证的场景:

#![allow(unused)]
fn main() {
#[derive(Debug, Clone)]
pub struct OAuthAuth {
    access_token: String,
    refresh_token: Option<String>,
    expires_at: Option<Instant>,
    token_type: String,
}

impl OAuthAuth {
    pub async fn initiate_device_flow(
        client_id: &str,
        scope: &str,
    ) -> CodexResult<DeviceFlowResponse> {
        let device_auth_url = "https://auth.openai.com/device/code";

        let params = [
            ("client_id", client_id),
            ("scope", scope),
        ];

        let response: DeviceFlowResponse = reqwest::Client::new()
            .post(device_auth_url)
            .form(&params)
            .send()
            .await?
            .json()
            .await?;

        Ok(response)
    }

    pub async fn poll_for_token(
        device_code: &str,
        client_id: &str,
        interval: Duration,
    ) -> CodexResult<OAuthAuth> {
        let token_url = "https://auth.openai.com/token";
        let client = reqwest::Client::new();

        loop {
            let params = [
                ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"),
                ("device_code", device_code),
                ("client_id", client_id),
            ];

            let response = client
                .post(token_url)
                .form(&params)
                .send()
                .await?;

            match response.status() {
                StatusCode::OK => {
                    let token_response: TokenResponse = response.json().await?;
                    return Ok(OAuthAuth::from_token_response(token_response));
                },
                StatusCode::BAD_REQUEST => {
                    let error: OAuthError = response.json().await?;
                    match error.error.as_str() {
                        "authorization_pending" => {
                            tokio::time::sleep(interval).await;
                            continue;
                        },
                        "slow_down" => {
                            tokio::time::sleep(interval * 2).await;
                            continue;
                        },
                        _ => return Err(CodexErr::AuthError(error.error)),
                    }
                },
                _ => return Err(CodexErr::AuthError("Token polling failed".to_string())),
            }
        }
    }
}
}

5.3 请求构建与模型选择

5.3.1 动态模型路由

Codex 支持根据任务复杂度、成本考虑等因素动态选择模型:

#![allow(unused)]
fn main() {
pub struct ModelRouter {
    available_models: Vec<ModelInfo>,
    routing_strategy: RoutingStrategy,
    cost_optimizer: CostOptimizer,
}

#[derive(Debug, Clone)]
pub enum RoutingStrategy {
    FastestFirst,    // 优先选择延迟最低的模型
    CostEfficient,   // 优先选择性价比最高的模型
    QualityFirst,    // 优先选择输出质量最高的模型
    LoadBalanced,    // 负载均衡
    Custom(Box<dyn Fn(&TaskContext) -> ModelId>), // 自定义策略
}

impl ModelRouter {
    pub fn select_model(&self, context: &TurnContext) -> CodexResult<ModelInfo> {
        let candidates = self.filter_capable_models(context)?;

        match &self.routing_strategy {
            RoutingStrategy::FastestFirst => {
                candidates
                    .iter()
                    .min_by_key(|model| model.avg_latency_ms)
                    .cloned()
                    .ok_or(CodexErr::NoSuitableModel)
            },
            RoutingStrategy::CostEfficient => {
                let optimal = self.cost_optimizer
                    .find_optimal_model(&candidates, context)?;
                Ok(optimal)
            },
            // ... 其他策略
        }
    }

    fn filter_capable_models(&self, context: &TurnContext) -> CodexResult<Vec<ModelInfo>> {
        self.available_models
            .iter()
            .filter(|model| {
                // 检查上下文窗口大小
                context.estimated_tokens() <= model.context_window &&
                // 检查功能支持(如 function calling)
                context.requires_function_calling() <= model.supports_functions &&
                // 检查多模态支持
                context.has_images() <= model.supports_vision
            })
            .cloned()
            .collect()
    }
}
}

5.3.2 请求构建流水线

每个 API 请求经过标准化的构建流程:

#![allow(unused)]
fn main() {
pub struct RequestBuilder {
    base_config: RequestConfig,
    prompt_builder: PromptBuilder,
    tool_serializer: ToolSerializer,
}

impl RequestBuilder {
    pub async fn build_completion_request(
        &self,
        context: &TurnContext,
        model_info: &ModelInfo,
    ) -> CodexResult<ApiRequest> {
        let mut request = ApiRequest::new(model_info.model_id.clone());

        // 1. 构建消息序列
        let messages = self.prompt_builder
            .build_message_sequence(context)
            .await?;
        request.set_messages(messages);

        // 2. 添加工具定义
        if context.requires_function_calling() {
            let tools = self.tool_serializer
                .serialize_available_tools(context.available_tools())
                .await?;
            request.set_tools(tools);
        }

        // 3. 设置生成参数
        request.set_temperature(self.base_config.temperature);
        request.set_max_tokens(self.calculate_max_tokens(context, model_info)?);
        request.set_stream(true); // 默认启用流式响应

        // 4. 添加提供商特定配置
        match model_info.provider_type {
            ProviderType::OpenAI => {
                self.apply_openai_specific_config(&mut request)?;
            },
            ProviderType::ResponsesApi => {
                self.apply_responses_api_config(&mut request, context)?;
            },
            // ... 其他提供商
        }

        Ok(request)
    }

    fn calculate_max_tokens(
        &self,
        context: &TurnContext,
        model_info: &ModelInfo,
    ) -> CodexResult<u32> {
        let input_tokens = context.estimated_tokens();
        let available_tokens = model_info.context_window.saturating_sub(input_tokens);

        // 保留一定的缓冲区,避免超出上下文窗口
        let buffer_tokens = (available_tokens as f32 * 0.1) as u32;
        let max_output = available_tokens.saturating_sub(buffer_tokens);

        // 限制在合理范围内
        Ok(max_output.min(model_info.max_output_tokens.unwrap_or(4096)))
    }
}
}

5.4 流式响应处理

5.4.1 双层流式架构

Codex 实现了两层流式处理架构:传输层流式 + 应用层流式:

#![allow(unused)]
fn main() {
pub struct StreamingResponse {
    transport_stream: Box<dyn Stream<Item = Result<Bytes, TransportError>>>,
    parser: ResponseParser,
    event_sender: mpsc::UnboundedSender<StreamEvent>,
}

#[derive(Debug, Clone)]
pub enum StreamEvent {
    TextDelta {
        delta: String,
        cumulative_text: String,
    },
    FunctionCall {
        name: String,
        arguments_delta: String,
        arguments_so_far: String,
    },
    ToolCallComplete {
        tool_call: ToolCall,
    },
    Usage {
        prompt_tokens: u32,
        completion_tokens: u32,
        total_tokens: u32,
    },
    Done,
    Error(StreamError),
}

impl StreamingResponse {
    pub async fn process_stream(mut self) -> CodexResult<()> {
        while let Some(chunk) = self.transport_stream.next().await {
            match chunk {
                Ok(bytes) => {
                    let events = self.parser.parse_chunk(&bytes)?;

                    for event in events {
                        if let Err(_) = self.event_sender.send(event) {
                            // 接收方已关闭,停止处理
                            break;
                        }
                    }
                },
                Err(transport_error) => {
                    let _ = self.event_sender.send(
                        StreamEvent::Error(StreamError::Transport(transport_error))
                    );
                    break;
                }
            }
        }

        let _ = self.event_sender.send(StreamEvent::Done);
        Ok(())
    }
}
}

5.4.2 SSE 与 WebSocket 双协议支持

根据提供商和场景,Codex 自动选择最适合的传输协议:

#![allow(unused)]
fn main() {
pub enum TransportType {
    ServerSentEvents,
    WebSocket,
    Http, // 非流式,用于简单请求
}

pub trait Transport: Send + Sync {
    async fn send_request(&mut self, request: ApiRequest) -> CodexResult<()>;
    async fn receive_stream(&mut self) -> CodexResult<StreamingResponse>;
    async fn close(&mut self) -> CodexResult<()>;
}

pub struct SseTransport {
    client: reqwest::Client,
    endpoint: String,
    headers: HeaderMap,
}

impl Transport for SseTransport {
    async fn send_request(&mut self, request: ApiRequest) -> CodexResult<()> {
        let response = self.client
            .post(&self.endpoint)
            .headers(self.headers.clone())
            .json(&request)
            .send()
            .await?;

        if !response.status().is_success() {
            return Err(CodexErr::HttpError(response.status()));
        }

        Ok(())
    }

    async fn receive_stream(&mut self) -> CodexResult<StreamingResponse> {
        // 实现 SSE 流式接收
        let stream = EventStreamReader::new(response.bytes_stream())
            .map(|result| {
                result
                    .map_err(|e| TransportError::EventStreamError(e))
                    .and_then(|event| {
                        match event.data.as_str() {
                            "[DONE]" => Ok(Bytes::new()), // 结束标记
                            data => Ok(Bytes::from(data.to_string())),
                        }
                    })
            });

        Ok(StreamingResponse::new(Box::new(stream)))
    }
}

pub struct WebSocketTransport {
    connection: Option<WebSocketConnection>,
    endpoint: String,
    auth_headers: HeaderMap,
}

impl Transport for WebSocketTransport {
    async fn send_request(&mut self, request: ApiRequest) -> CodexResult<()> {
        let ws_conn = match &mut self.connection {
            Some(conn) if !conn.is_closed() => conn,
            _ => {
                // 建立新的 WebSocket 连接
                let conn = self.establish_websocket_connection().await?;
                self.connection = Some(conn);
                self.connection.as_mut().unwrap()
            }
        };

        let message = serde_json::to_string(&request)?;
        ws_conn.send(Message::Text(message)).await?;

        Ok(())
    }

    async fn receive_stream(&mut self) -> CodexResult<StreamingResponse> {
        let conn = self.connection.as_mut()
            .ok_or(CodexErr::WebSocketNotConnected)?;

        let stream = conn.message_stream()
            .map(|result| {
                result
                    .map_err(|e| TransportError::WebSocketError(e))
                    .and_then(|message| {
                        match message {
                            Message::Text(text) => Ok(Bytes::from(text)),
                            Message::Binary(data) => Ok(Bytes::from(data)),
                            Message::Close(_) => Err(TransportError::ConnectionClosed),
                            _ => Ok(Bytes::new()),
                        }
                    })
            });

        Ok(StreamingResponse::new(Box::new(stream)))
    }
}
}

5.4.3 响应解析器

不同提供商返回的格式需要统一解析:

#![allow(unused)]
fn main() {
pub struct ResponseParser {
    provider_type: ProviderType,
    accumulated_content: String,
    current_tool_calls: HashMap<String, PartialToolCall>,
}

impl ResponseParser {
    pub fn parse_chunk(&mut self, chunk: &[u8]) -> CodexResult<Vec<StreamEvent>> {
        let text = std::str::from_utf8(chunk)?;

        match self.provider_type {
            ProviderType::OpenAI => self.parse_openai_chunk(text),
            ProviderType::ResponsesApi => self.parse_responses_api_chunk(text),
            _ => self.parse_generic_chunk(text),
        }
    }

    fn parse_openai_chunk(&mut self, chunk: &str) -> CodexResult<Vec<StreamEvent>> {
        let mut events = Vec::new();

        for line in chunk.lines() {
            if !line.starts_with("data: ") {
                continue;
            }

            let json_data = &line[6..]; // 跳过 "data: "

            if json_data == "[DONE]" {
                events.push(StreamEvent::Done);
                break;
            }

            let delta: serde_json::Value = serde_json::from_str(json_data)?;

            if let Some(choices) = delta["choices"].as_array() {
                for choice in choices {
                    if let Some(delta_obj) = choice["delta"].as_object() {
                        // 处理文本增量
                        if let Some(content) = delta_obj["content"].as_str() {
                            self.accumulated_content.push_str(content);
                            events.push(StreamEvent::TextDelta {
                                delta: content.to_string(),
                                cumulative_text: self.accumulated_content.clone(),
                            });
                        }

                        // 处理工具调用
                        if let Some(tool_calls) = delta_obj["tool_calls"].as_array() {
                            for tool_call in tool_calls {
                                let parsed_events = self.parse_tool_call_delta(tool_call)?;
                                events.extend(parsed_events);
                            }
                        }
                    }

                    // 处理使用统计
                    if let Some(usage) = choice["usage"].as_object() {
                        events.push(self.parse_usage_info(usage)?);
                    }
                }
            }
        }

        Ok(events)
    }

    fn parse_tool_call_delta(&mut self, tool_call: &serde_json::Value) -> CodexResult<Vec<StreamEvent>> {
        let mut events = Vec::new();

        let index = tool_call["index"].as_u64().unwrap_or(0) as usize;
        let call_id = tool_call["id"].as_str().unwrap_or("").to_string();

        // 获取或创建部分工具调用状态
        let partial_call = self.current_tool_calls
            .entry(call_id.clone())
            .or_insert_with(|| PartialToolCall::new(call_id.clone()));

        // 处理函数名称
        if let Some(function) = tool_call["function"].as_object() {
            if let Some(name) = function["name"].as_str() {
                partial_call.name = Some(name.to_string());
            }

            if let Some(arguments_delta) = function["arguments"].as_str() {
                partial_call.arguments.push_str(arguments_delta);

                events.push(StreamEvent::FunctionCall {
                    name: partial_call.name.clone().unwrap_or_default(),
                    arguments_delta: arguments_delta.to_string(),
                    arguments_so_far: partial_call.arguments.clone(),
                });
            }
        }

        // 检查工具调用是否完成
        if let Some(finish_reason) = tool_call["finish_reason"].as_str() {
            if finish_reason == "tool_calls" || finish_reason == "stop" {
                if let Some(completed_call) = self.current_tool_calls.remove(&call_id) {
                    events.push(StreamEvent::ToolCallComplete {
                        tool_call: completed_call.into_tool_call()?,
                    });
                }
            }
        }

        Ok(events)
    }
}
}

5.5 重试机制与错误恢复

5.5.1 分层重试策略

Codex 实现了多层级的重试机制:

#![allow(unused)]
fn main() {
#[derive(Debug, Clone)]
pub struct RetryConfig {
    max_attempts: usize,
    base_delay: Duration,
    max_delay: Duration,
    backoff_strategy: BackoffStrategy,
    retryable_errors: HashSet<ErrorType>,
}

#[derive(Debug, Clone)]
pub enum BackoffStrategy {
    Fixed,
    Linear,
    Exponential { multiplier: f32 },
    Jittered { base_multiplier: f32, jitter_ratio: f32 },
}

pub struct RetryableApiClient {
    inner_client: Box<dyn Transport>,
    retry_config: RetryConfig,
    circuit_breaker: CircuitBreaker,
    metrics: Arc<ApiMetrics>,
}

impl RetryableApiClient {
    pub async fn execute_with_retry<F, T>(&self, operation: F) -> CodexResult<T>
    where
        F: Fn() -> BoxFuture<'_, CodexResult<T>>,
        T: Clone,
    {
        let mut attempt = 0;
        let mut last_error = None;

        while attempt < self.retry_config.max_attempts {
            // 检查熔断器状态
            match self.circuit_breaker.check_state() {
                CircuitState::Open => {
                    return Err(CodexErr::CircuitBreakerOpen);
                },
                CircuitState::HalfOpen if attempt > 0 => {
                    // 半开状态下,只允许第一次尝试
                    return Err(last_error.unwrap_or(CodexErr::MaxRetriesExceeded));
                },
                _ => {}
            }

            let start_time = Instant::now();
            let result = operation().await;
            let duration = start_time.elapsed();

            self.metrics.record_attempt(attempt, duration, result.is_ok());

            match result {
                Ok(value) => {
                    // 成功时重置熔断器
                    self.circuit_breaker.record_success();
                    return Ok(value);
                },
                Err(error) => {
                    attempt += 1;
                    last_error = Some(error.clone());

                    // 检查错误是否可重试
                    if !self.is_retryable_error(&error) {
                        self.circuit_breaker.record_failure();
                        return Err(error);
                    }

                    // 计算退避延迟
                    if attempt < self.retry_config.max_attempts {
                        let delay = self.calculate_backoff_delay(attempt);
                        tokio::time::sleep(delay).await;
                    }
                }
            }
        }

        // 所有重试都失败
        self.circuit_breaker.record_failure();
        Err(last_error.unwrap_or(CodexErr::MaxRetriesExceeded))
    }

    fn calculate_backoff_delay(&self, attempt: usize) -> Duration {
        let base_delay = self.retry_config.base_delay;

        let delay = match &self.retry_config.backoff_strategy {
            BackoffStrategy::Fixed => base_delay,

            BackoffStrategy::Linear => {
                Duration::from_millis(base_delay.as_millis() as u64 * attempt as u64)
            },

            BackoffStrategy::Exponential { multiplier } => {
                let multiplied_ms = base_delay.as_millis() as f32 * multiplier.powi(attempt as i32);
                Duration::from_millis(multiplied_ms as u64)
            },

            BackoffStrategy::Jittered { base_multiplier, jitter_ratio } => {
                let base_ms = base_delay.as_millis() as f32 * base_multiplier.powi(attempt as i32);
                let jitter = base_ms * jitter_ratio * rand::random::<f32>();
                Duration::from_millis((base_ms + jitter) as u64)
            },
        };

        delay.min(self.retry_config.max_delay)
    }

    fn is_retryable_error(&self, error: &CodexErr) -> bool {
        match error {
            CodexErr::Network(_) => true,
            CodexErr::RateLimit { .. } => true,
            CodexErr::ServerError(status) if status.is_server_error() => true,
            CodexErr::Timeout => true,
            CodexErr::TooManyRequests { .. } => true,
            _ => false,
        }
    }
}
}

5.5.2 Rate Limiting 实现

#![allow(unused)]
fn main() {
use tokio::sync::Semaphore;
use std::sync::Arc;
use std::time::{Duration, Instant};

pub struct RateLimiter {
    // Token bucket for requests per second
    request_semaphore: Arc<Semaphore>,
    // Token bucket for tokens per minute
    token_semaphore: Arc<Semaphore>,
    // Sliding window for tracking request history
    request_history: Arc<Mutex<VecDeque<Instant>>>,

    config: RateLimitConfig,
}

#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    requests_per_second: u32,
    tokens_per_minute: u32,
    burst_allowance: u32,
    window_size: Duration,
}

impl RateLimiter {
    pub async fn acquire_request_permit(&self, estimated_tokens: u32) -> CodexResult<RateLimitPermit> {
        // 1. 检查请求频率限制
        let _request_permit = self.request_semaphore.acquire().await?;

        // 2. 检查 Token 使用限制
        let _token_permit = self.token_semaphore.acquire_many(estimated_tokens).await?;

        // 3. 更新请求历史
        {
            let mut history = self.request_history.lock().await;
            let now = Instant::now();

            // 清理过期记录
            while let Some(&front_time) = history.front() {
                if now.duration_since(front_time) > self.config.window_size {
                    history.pop_front();
                } else {
                    break;
                }
            }

            history.push_back(now);

            // 检查突发请求限制
            if history.len() > self.config.burst_allowance as usize {
                return Err(CodexErr::RateLimitExceeded {
                    retry_after: Some(Duration::from_secs(1)),
                });
            }
        }

        Ok(RateLimitPermit {
            _request_permit,
            _token_permit,
            acquired_tokens: estimated_tokens,
        })
    }

    // 后台 Token 恢复任务
    async fn token_recovery_task(&self) {
        let mut interval = tokio::time::interval(Duration::from_secs(60));

        loop {
            interval.tick().await;

            // 每分钟恢复可用 Token 数量
            let permits_to_add = self.config.tokens_per_minute - self.token_semaphore.available_permits() as u32;
            if permits_to_add > 0 {
                self.token_semaphore.add_permits(permits_to_add as usize);
            }
        }
    }
}

pub struct RateLimitPermit {
    _request_permit: SemaphorePermit<'static>,
    _token_permit: SemaphorePermit<'static>,
    acquired_tokens: u32,
}

impl Drop for RateLimitPermit {
    fn drop(&mut self) {
        // Permits 会自动释放
        tracing::debug!("Released rate limit permit for {} tokens", self.acquired_tokens);
    }
}
}

5.6 连接池与会话管理

5.6.1 WebSocket 连接复用

对于支持 WebSocket 的 API(如 Responses API),Codex 实现了连接池来提高性能:

#![allow(unused)]
fn main() {
pub struct ConnectionPool {
    pools: Arc<RwLock<HashMap<String, Pool<WebSocketConnection>>>>,
    config: ConnectionPoolConfig,
}

#[derive(Debug, Clone)]
pub struct ConnectionPoolConfig {
    max_connections_per_endpoint: usize,
    idle_timeout: Duration,
    max_lifetime: Duration,
    health_check_interval: Duration,
}

pub struct Pool<T> {
    connections: VecDeque<PooledConnection<T>>,
    active_count: usize,
    max_size: usize,
}

struct PooledConnection<T> {
    connection: T,
    created_at: Instant,
    last_used: Instant,
    is_healthy: bool,
}

impl ConnectionPool {
    pub async fn get_connection(&self, endpoint: &str) -> CodexResult<PooledWebSocketConnection> {
        let pools = self.pools.read().await;

        if let Some(pool) = pools.get(endpoint) {
            if let Some(conn) = pool.try_get_connection() {
                if conn.is_healthy() {
                    return Ok(PooledWebSocketConnection::new(conn, endpoint.to_string()));
                }
            }
        }

        drop(pools);

        // 没有可用连接,创建新连接
        self.create_new_connection(endpoint).await
    }

    async fn create_new_connection(&self, endpoint: &str) -> CodexResult<PooledWebSocketConnection> {
        let mut pools = self.pools.write().await;
        let pool = pools.entry(endpoint.to_string())
            .or_insert_with(|| Pool::new(self.config.max_connections_per_endpoint));

        if pool.active_count >= pool.max_size {
            return Err(CodexErr::ConnectionPoolExhausted);
        }

        let ws_connection = self.establish_websocket_connection(endpoint).await?;
        pool.active_count += 1;

        Ok(PooledWebSocketConnection::new(ws_connection, endpoint.to_string()))
    }

    async fn establish_websocket_connection(&self, endpoint: &str) -> CodexResult<WebSocketConnection> {
        let (ws_stream, _) = tokio_tungstenite::connect_async(endpoint).await?;

        Ok(WebSocketConnection {
            stream: ws_stream,
            created_at: Instant::now(),
            last_ping: Instant::now(),
        })
    }

    // 后台健康检查任务
    pub async fn health_check_task(&self) {
        let mut interval = tokio::time::interval(self.config.health_check_interval);

        loop {
            interval.tick().await;
            self.check_and_cleanup_connections().await;
        }
    }

    async fn check_and_cleanup_connections(&self) {
        let mut pools = self.pools.write().await;

        for (endpoint, pool) in pools.iter_mut() {
            let now = Instant::now();

            pool.connections.retain(|conn| {
                let is_expired = now.duration_since(conn.created_at) > self.config.max_lifetime ||
                                now.duration_since(conn.last_used) > self.config.idle_timeout;

                if is_expired {
                    pool.active_count = pool.active_count.saturating_sub(1);
                    false
                } else {
                    true
                }
            });
        }

        // 清理空的池
        pools.retain(|_, pool| !pool.connections.is_empty() || pool.active_count > 0);
    }
}
}

5.6.2 Turn 级会话管理

每个 Turn 创建独立的会话来处理 WebSocket 连接的生命周期:

#![allow(unused)]
fn main() {
impl ModelClientSession {
    pub async fn new(client: Arc<ModelClient>, turn_context: TurnContext) -> CodexResult<Self> {
        Ok(Self {
            client,
            turn_context,
            websocket_connection: None,
            turn_state_token: None,
        })
    }

    pub async fn send_streaming_request(
        &mut self,
        request: ApiRequest,
    ) -> CodexResult<StreamingResponse> {
        match self.client.provider_config.provider_type {
            ProviderType::ResponsesApi => {
                self.send_websocket_request(request).await
            },
            _ => {
                self.send_http_streaming_request(request).await
            }
        }
    }

    async fn send_websocket_request(&mut self, request: ApiRequest) -> CodexResult<StreamingResponse> {
        // 获取或建立 WebSocket 连接
        let connection = match &mut self.websocket_connection {
            Some(conn) if conn.is_alive() => conn,
            _ => {
                let new_conn = self.client.connection_pool
                    .get_connection(&self.client.provider_config.endpoint_url)
                    .await?;
                self.websocket_connection = Some(new_conn);
                self.websocket_connection.as_mut().unwrap()
            }
        };

        // 如果有前一个响应的 ID,用于粘性路由
        let mut ws_request = ResponsesWsRequest::from(request);
        if let Some(prev_response_id) = &self.turn_state_token {
            ws_request.set_previous_response_id(prev_response_id.clone());
        }

        // 发送请求
        connection.send_request(ws_request).await?;

        // 返回流式响应
        let stream = connection.receive_stream().await?;

        Ok(stream)
    }
}

pub struct PooledWebSocketConnection {
    connection: WebSocketConnection,
    endpoint: String,
    return_to_pool: Option<oneshot::Sender<WebSocketConnection>>,
}

impl Drop for PooledWebSocketConnection {
    fn drop(&mut self) {
        if let Some(sender) = self.return_to_pool.take() {
            // 将连接返回到池中
            let _ = sender.send(self.connection.clone());
        }
    }
}
}

5.7 性能优化与监控

5.7.1 请求预热机制

为了减少 WebSocket 建立连接的延迟,Codex 实现了连接预热:

#![allow(unused)]
fn main() {
pub struct ConnectionPrewarmer {
    client: Arc<ModelClient>,
    prewarming_queue: Arc<Mutex<VecDeque<PrewarmTask>>>,
    worker_handles: Vec<JoinHandle<()>>,
}

struct PrewarmTask {
    endpoint: String,
    priority: PrewarmPriority,
    created_at: Instant,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum PrewarmPriority {
    Low = 0,
    Normal = 1,
    High = 2,
    Urgent = 3,
}

impl ConnectionPrewarmer {
    pub fn new(client: Arc<ModelClient>, worker_count: usize) -> Self {
        let prewarming_queue = Arc::new(Mutex::new(VecDeque::new()));
        let mut worker_handles = Vec::new();

        // 启动预热工作线程
        for _ in 0..worker_count {
            let queue = prewarming_queue.clone();
            let client = client.clone();

            let handle = tokio::spawn(async move {
                Self::prewarm_worker(queue, client).await;
            });

            worker_handles.push(handle);
        }

        Self {
            client,
            prewarming_queue,
            worker_handles,
        }
    }

    pub async fn schedule_prewarm(&self, endpoint: &str, priority: PrewarmPriority) {
        let task = PrewarmTask {
            endpoint: endpoint.to_string(),
            priority,
            created_at: Instant::now(),
        };

        let mut queue = self.prewarming_queue.lock().await;

        // 插入队列,按优先级排序
        let insert_pos = queue
            .iter()
            .position(|existing| existing.priority < priority)
            .unwrap_or(queue.len());

        queue.insert(insert_pos, task);
    }

    async fn prewarm_worker(
        queue: Arc<Mutex<VecDeque<PrewarmTask>>>,
        client: Arc<ModelClient>,
    ) {
        loop {
            let task = {
                let mut queue = queue.lock().await;
                queue.pop_front()
            };

            if let Some(task) = task {
                // 执行预热:发送一个 generate=false 的请求
                match Self::execute_prewarm(&client, &task.endpoint).await {
                    Ok(_) => {
                        tracing::debug!("Successfully prewarmed connection to {}", task.endpoint);
                    },
                    Err(err) => {
                        tracing::warn!("Failed to prewarm connection to {}: {}", task.endpoint, err);
                    }
                }
            } else {
                // 队列为空,短暂休眠
                tokio::time::sleep(Duration::from_millis(100)).await;
            }
        }
    }

    async fn execute_prewarm(client: &ModelClient, endpoint: &str) -> CodexResult<()> {
        let prewarm_request = ApiRequest {
            model: "gpt-4".to_string(), // 使用默认模型
            messages: vec![], // 空消息
            generate: false, // 关键:不生成响应,只建立连接
            stream: true,
        };

        let mut session = ModelClientSession::new(client.clone(), TurnContext::empty()).await?;
        let _response = session.send_streaming_request(prewarm_request).await?;

        // 立即关闭流,我们只是想建立连接
        Ok(())
    }
}
}

5.7.2 性能指标收集

#![allow(unused)]
fn main() {
#[derive(Debug, Clone)]
pub struct ApiMetrics {
    request_count: Arc<AtomicU64>,
    success_count: Arc<AtomicU64>,
    error_count: Arc<AtomicU64>,
    total_latency: Arc<AtomicU64>, // 毫秒
    token_usage: Arc<AtomicU64>,

    latency_histogram: Arc<Mutex<Histogram>>,
    error_breakdown: Arc<Mutex<HashMap<String, u64>>>,
}

impl ApiMetrics {
    pub fn record_request(&self, duration: Duration, success: bool, tokens: Option<u32>) {
        self.request_count.fetch_add(1, Ordering::Relaxed);

        if success {
            self.success_count.fetch_add(1, Ordering::Relaxed);
        } else {
            self.error_count.fetch_add(1, Ordering::Relaxed);
        }

        let latency_ms = duration.as_millis() as u64;
        self.total_latency.fetch_add(latency_ms, Ordering::Relaxed);

        if let Some(token_count) = tokens {
            self.token_usage.fetch_add(token_count as u64, Ordering::Relaxed);
        }

        // 更新延迟直方图
        {
            let mut histogram = self.latency_histogram.lock().unwrap();
            histogram.record(latency_ms);
        }
    }

    pub fn record_error(&self, error_type: &str) {
        let mut breakdown = self.error_breakdown.lock().unwrap();
        *breakdown.entry(error_type.to_string()).or_insert(0) += 1;
    }

    pub fn get_summary(&self) -> ApiMetricsSummary {
        let request_count = self.request_count.load(Ordering::Relaxed);
        let success_count = self.success_count.load(Ordering::Relaxed);
        let error_count = self.error_count.load(Ordering::Relaxed);
        let total_latency = self.total_latency.load(Ordering::Relaxed);
        let token_usage = self.token_usage.load(Ordering::Relaxed);

        let success_rate = if request_count > 0 {
            success_count as f64 / request_count as f64
        } else {
            0.0
        };

        let average_latency = if request_count > 0 {
            total_latency as f64 / request_count as f64
        } else {
            0.0
        };

        ApiMetricsSummary {
            request_count,
            success_rate,
            average_latency_ms: average_latency,
            total_tokens_used: token_usage,
            error_breakdown: self.error_breakdown.lock().unwrap().clone(),
        }
    }
}

struct Histogram {
    buckets: Vec<(u64, u64)>, // (upper_bound, count)
}

impl Histogram {
    fn new() -> Self {
        Self {
            buckets: vec![
                (50, 0),    // 0-50ms
                (100, 0),   // 51-100ms
                (200, 0),   // 101-200ms
                (500, 0),   // 201-500ms
                (1000, 0),  // 501-1000ms
                (2000, 0),  // 1001-2000ms
                (5000, 0),  // 2001-5000ms
                (u64::MAX, 0), // >5000ms
            ],
        }
    }

    fn record(&mut self, value: u64) {
        for (upper_bound, count) in &mut self.buckets {
            if value <= *upper_bound {
                *count += 1;
                break;
            }
        }
    }
}
}

5.8 总结与设计洞察

5.8.1 核心设计原则

OpenAI Codex CLI 的 API Client 体现了以下设计原则:

  1. Provider Agnostic:统一接口屏蔽提供商差异
  2. Resilience First:多层重试和熔断保证可靠性
  3. Performance Optimized:连接复用和预热提升性能
  4. Observable:丰富的指标和日志便于运维
  5. Secure by Default:内置认证和权限控制

5.8.2 关键技术选择

技术选择理由权衡
双传输协议WebSocket 低延迟,HTTP 兼容性好增加复杂性
连接池减少连接建立开销内存占用增加
流式解析实时响应,用户体验好状态管理复杂
熔断器模式故障隔离,快速失败可能误判暂时性故障

5.8.3 架构对比

维度Codex CLIClaude Code优势分析
连接管理连接池 + 预热简单连接Codex 性能更优
错误处理分层重试 + 熔断基础重试Codex 可靠性更高
协议支持HTTP + WebSocket仅 HTTPCodex 延迟更低
监控能力详细指标基础日志Codex 可观测性更强

5.8.4 速查表

组件文件路径核心功能关键接口
API 客户端client.rs模型 API 调用ModelClient::stream()
认证管理auth.rs多种认证方式AuthManager::get_auth()
传输层transport.rsHTTP/WebSocketTransport::send()
响应解析parser.rs流式响应解析ResponseParser::parse()
连接池connection_pool.rs连接复用ConnectionPool::get()
重试机制retry.rs错误恢复RetryableClient::execute()

Codex 的 API Client 是一个经过精心设计的通信引擎,它在保证高性能的同时提供了出色的可靠性和可扩展性。这个架构为构建生产级 AI Agent 系统提供了重要的参考价值。