Skip to main content

llm/providers/codex/
provider.rs

1use super::oauth::CodexTokenManager;
2use crate::provider::{LlmResponseStream, StreamingModelProvider, get_context_window, stream_from};
3use crate::providers::openai_responses::mappers::{ResponsesRequestPolicy, build_wire_request};
4use crate::providers::openai_responses::streaming::process_response_stream;
5use crate::providers::openai_responses::transport::{ResponsesEventStream, decode_response_sse};
6use crate::{Context, LlmError, Result};
7use aether_auth::OAuthCredentialStorage;
8use reqwest::header::{ACCEPT, AUTHORIZATION, HeaderMap, HeaderValue};
9use std::sync::Arc;
10use tracing::debug;
11
12const CODEX_API_BASE: &str = "https://chatgpt.com/backend-api/codex";
13const CODEX_CLIENT_VERSION: &str = "0.144.0";
14
15#[derive(Clone)]
16pub struct CodexProvider {
17    base_url: String,
18    client: reqwest::Client,
19    model: String,
20    token_manager: Arc<CodexTokenManager>,
21}
22
23impl CodexProvider {
24    pub fn new(store: Arc<dyn OAuthCredentialStorage>) -> Self {
25        let token_manager = CodexTokenManager::new(store, super::PROVIDER_ID);
26        Self {
27            base_url: CODEX_API_BASE.to_string(),
28            client: reqwest::Client::new(),
29            model: "gpt-5.5".to_string(),
30            token_manager: Arc::new(token_manager),
31        }
32    }
33
34    pub fn with_connection(mut self, connection: crate::ProviderConnectionConfig) -> Self {
35        if let Some(base_url) = connection.base_url {
36            self.base_url = base_url.trim_end_matches('/').to_string();
37        }
38        self
39    }
40
41    pub fn with_model(mut self, model: &str) -> Self {
42        self.model = model.to_string();
43        self
44    }
45
46    fn build_wire_request(&self, context: &Context) -> Result<serde_json::Value> {
47        build_wire_request(&self.model, context, &ResponsesRequestPolicy::codex())
48    }
49
50    async fn build_headers(&self) -> Result<HeaderMap> {
51        let (access_token, account_id) = self.token_manager.get_valid_token().await?;
52
53        let mut headers = HeaderMap::new();
54        headers.insert(
55            AUTHORIZATION,
56            HeaderValue::from_str(&format!("Bearer {access_token}"))
57                .map_err(|e| LlmError::InvalidApiKey(e.to_string()))?,
58        );
59        headers.insert(
60            "chatgpt-account-id",
61            HeaderValue::from_str(&account_id).map_err(|e| LlmError::InvalidApiKey(e.to_string()))?,
62        );
63        headers.insert(ACCEPT, HeaderValue::from_static("text/event-stream"));
64        headers.insert("originator", HeaderValue::from_static("codex_cli_rs"));
65        headers.insert("version", HeaderValue::from_static(CODEX_CLIENT_VERSION));
66
67        Ok(headers)
68    }
69
70    /// Send the request and return a stream of SSE lines parsed into typed events.
71    ///
72    /// Uses manual SSE parsing because the Codex API does not return a
73    /// `Content-Type: text/event-stream` header, which `reqwest_eventsource`
74    /// (used by `async-openai`'s `create_stream`) requires.
75    async fn send_request(&self, request: serde_json::Value, headers: HeaderMap) -> Result<ResponsesEventStream> {
76        let url = format!("{}/responses", self.base_url);
77
78        debug!("Sending request to Codex API: {url}");
79        debug!(
80            "Codex request body: {}",
81            serde_json::to_string(&request).unwrap_or_else(|_| "<failed to serialize>".to_string())
82        );
83
84        let response = self.client.post(&url).headers(headers).json(&request).send().await?;
85
86        if !response.status().is_success() {
87            let status = response.status();
88            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
89
90            if matches!(status.as_u16(), 401 | 403) {
91                self.token_manager.clear_cache().await;
92            }
93
94            let message = format!("Codex API request failed with status {status}: {error_text}");
95            return Err(match status.as_u16() {
96                429 => LlmError::RateLimited(message),
97                s if (500..600).contains(&s) => LlmError::ServerError { status: Some(s), message },
98                _ => LlmError::ApiError(message),
99            });
100        }
101
102        Ok(decode_response_sse(response))
103    }
104}
105
106impl StreamingModelProvider for CodexProvider {
107    fn model(&self) -> Option<crate::LlmModel> {
108        format!("{}:{}", super::PROVIDER_ID, self.model).parse().ok()
109    }
110
111    fn context_window(&self) -> Option<u32> {
112        get_context_window(super::PROVIDER_ID, &self.model)
113    }
114
115    fn stream_response(&self, context: &Context) -> LlmResponseStream {
116        let provider = self.clone();
117        let context = context.clone();
118
119        stream_from(
120            async move {
121                let headers = provider.build_headers().await?;
122                let request = provider.build_wire_request(&context)?;
123                provider.send_request(request, headers).await
124            },
125            process_response_stream,
126        )
127    }
128
129    fn display_name(&self) -> String {
130        format!("Codex ({})", self.model)
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use crate::ChatMessage;
138    use crate::ToolDefinition;
139    use crate::providers::test_capture_server::CaptureServer;
140    use aether_auth::{FakeOAuthCredentialStore, OAuthCredential};
141    use base64::Engine;
142    use base64::engine::general_purpose::URL_SAFE_NO_PAD;
143    use futures::StreamExt;
144
145    #[test]
146    fn context_window_uses_codex_subscription_limit() {
147        let provider = create_test_provider();
148        assert_eq!(provider.context_window(), Some(272_000));
149    }
150
151    #[test]
152    fn display_name_includes_model() {
153        let provider = create_test_provider();
154        assert_eq!(provider.display_name(), "Codex (gpt-5.5)");
155    }
156
157    #[tokio::test]
158    async fn stream_response_sends_supported_protocol_version_for_gpt_5_6_luna() {
159        let mut server = CaptureServer::start().await;
160        let provider = server_backed_provider(&server).with_model("gpt-5.6-luna");
161        let mut context = Context::new(
162            vec![ChatMessage::system("You are helpful"), ChatMessage::user("Think harder")],
163            vec![ToolDefinition::new(
164                "bash",
165                "Run a command",
166                serde_json::from_str(r#"{"type": "object", "properties": {"cmd": {"type": "string"}}}"#).unwrap(),
167            )],
168        );
169        context.set_reasoning_effort(Some(crate::ReasoningEffort::Max));
170        context.set_prompt_cache_key(Some("session-abc".to_string()));
171
172        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
173        let captured = server.captured().await;
174
175        assert!(responses.iter().all(Result::is_ok), "{responses:?}");
176        assert_eq!(captured.body["reasoning"]["effort"], "max");
177        assert!(captured.body["reasoning"].get("context").is_none());
178        assert_eq!(captured.body["model"], "gpt-5.6-luna");
179        assert_eq!(captured.body["instructions"], "You are helpful");
180        assert_eq!(captured.body["tools"].as_array().unwrap().len(), 1);
181        assert!(captured.body.get("parallel_tool_calls").is_none());
182        assert_eq!(captured.body["input"][0]["role"], "user");
183        assert_eq!(captured.body["prompt_cache_key"], "session-abc");
184        assert_eq!(captured.body["store"], false);
185        assert_eq!(captured.body["stream"], true);
186        assert_eq!(captured.headers["chatgpt-account-id"], "account-1");
187        assert_eq!(captured.headers["version"], "0.144.0");
188        assert_eq!(captured.headers["accept"], "text/event-stream");
189        assert!(captured.headers.get("x-openai-internal-codex-responses-lite").is_none());
190        assert!(captured.headers.get("OpenAI-Beta").is_none());
191    }
192
193    #[tokio::test]
194    async fn stream_response_defaults_to_medium_effort_on_the_wire() {
195        let mut server = CaptureServer::start().await;
196        let provider = server_backed_provider(&server);
197        let context = Context::new(vec![ChatMessage::user("Hello")], vec![]);
198
199        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
200        let captured = server.captured().await;
201
202        assert!(responses.iter().all(Result::is_ok), "{responses:?}");
203        assert_eq!(captured.body["reasoning"]["effort"], "medium");
204    }
205
206    fn server_backed_provider(server: &CaptureServer) -> CodexProvider {
207        let credential = OAuthCredential {
208            client_id: "test".to_string(),
209            access_token: test_jwt("account-1"),
210            refresh_token: None,
211            expires_at: Some(u64::MAX),
212            granted_scopes: Vec::new(),
213        };
214        let store: Arc<dyn OAuthCredentialStorage> =
215            Arc::new(FakeOAuthCredentialStore::new().with_credential("codex", credential));
216        CodexProvider::new(store).with_connection(crate::ProviderConnectionConfig {
217            base_url: Some(server.base_url.clone()),
218            ..Default::default()
219        })
220    }
221
222    fn test_jwt(account_id: &str) -> String {
223        let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none"}"#);
224        let payload = URL_SAFE_NO_PAD
225            .encode(serde_json::json!({"https://api.openai.com/auth": {"chatgpt_account_id": account_id}}).to_string());
226        format!("{header}.{payload}.signature")
227    }
228
229    fn create_test_provider() -> CodexProvider {
230        let store: Arc<dyn OAuthCredentialStorage> = Arc::new(FakeOAuthCredentialStore::new());
231        CodexProvider::new(store).with_model("gpt-5.5")
232    }
233}