llm/providers/codex/
provider.rs

1use super::mappers::{map_messages, map_tools};
2use super::oauth::CodexTokenManager;
3use super::streaming::{CodexStreamEvent, process_response_stream};
4use crate::provider::{LlmResponseStream, StreamingModelProvider, get_context_window};
5use crate::{Context, LlmError, Result};
6use aether_auth::OAuthCredentialStorage;
7use async_openai::types::responses::{
8    CreateResponse, IncludeEnum, InputParam, Reasoning, ReasoningSummary, ResponseTextParam,
9    TextResponseFormatConfiguration, Verbosity,
10};
11use eventsource_stream::Eventsource;
12use futures::StreamExt;
13use reqwest::header::{ACCEPT, AUTHORIZATION, HeaderMap, HeaderValue};
14use serde_json::Value;
15use std::sync::Arc;
16use tracing::debug;
17
18const CODEX_API_BASE: &str = "https://chatgpt.com/backend-api/codex";
19const CODEX_CLIENT_VERSION: &str = "0.144.0";
20
21#[derive(Clone)]
22pub struct CodexProvider {
23    base_url: String,
24    client: reqwest::Client,
25    model: String,
26    token_manager: Arc<CodexTokenManager>,
27}
28
29impl CodexProvider {
30    pub fn new(store: Arc<dyn OAuthCredentialStorage>) -> Self {
31        let token_manager = CodexTokenManager::new(store, super::PROVIDER_ID);
32        Self {
33            base_url: CODEX_API_BASE.to_string(),
34            client: reqwest::Client::new(),
35            model: "gpt-5.5".to_string(),
36            token_manager: Arc::new(token_manager),
37        }
38    }
39
40    pub fn with_connection(mut self, connection: crate::ProviderConnectionConfig) -> Self {
41        if let Some(base_url) = connection.base_url {
42            self.base_url = base_url.trim_end_matches('/').to_string();
43        }
44        self
45    }
46
47    pub fn with_model(mut self, model: &str) -> Self {
48        self.model = model.to_string();
49        self
50    }
51
52    fn build_typed_request(&self, context: &Context) -> Result<CreateResponse> {
53        let (system_prompt, input) = map_messages(context.messages())?;
54        let tools = if context.tools().is_empty() { None } else { Some(map_tools(context.tools())?) };
55
56        Ok(CreateResponse {
57            model: Some(self.model.clone()),
58            input: InputParam::Items(input),
59            instructions: system_prompt,
60            tools,
61            store: Some(false),
62            stream: Some(true),
63            reasoning: Some(Reasoning { effort: None, summary: Some(ReasoningSummary::Auto) }),
64            include: Some(vec![IncludeEnum::ReasoningEncryptedContent]),
65            text: Some(ResponseTextParam {
66                format: TextResponseFormatConfiguration::Text,
67                verbosity: Some(Verbosity::Medium),
68            }),
69            prompt_cache_key: context.prompt_cache_key().map(String::from),
70            ..Default::default()
71        })
72    }
73
74    /// The typed `async-openai` request cannot encode every effort the Codex
75    /// API accepts (e.g. `max`), so the effort is written onto the serialized
76    /// wire body instead.
77    fn build_wire_request(&self, context: &Context) -> Result<Value> {
78        let mut body = serde_json::to_value(self.build_typed_request(context)?)?;
79        let effort = context.reasoning_effort().map_or("medium", crate::ReasoningEffort::as_str);
80        body["reasoning"]["effort"] = effort.into();
81        Ok(body)
82    }
83
84    async fn build_headers(&self) -> Result<HeaderMap> {
85        let (access_token, account_id) = self.token_manager.get_valid_token().await?;
86
87        let mut headers = HeaderMap::new();
88        headers.insert(
89            AUTHORIZATION,
90            HeaderValue::from_str(&format!("Bearer {access_token}"))
91                .map_err(|e| LlmError::InvalidApiKey(e.to_string()))?,
92        );
93        headers.insert(
94            "chatgpt-account-id",
95            HeaderValue::from_str(&account_id).map_err(|e| LlmError::InvalidApiKey(e.to_string()))?,
96        );
97        headers.insert(ACCEPT, HeaderValue::from_static("text/event-stream"));
98        headers.insert("originator", HeaderValue::from_static("codex_cli_rs"));
99        headers.insert("version", HeaderValue::from_static(CODEX_CLIENT_VERSION));
100
101        Ok(headers)
102    }
103
104    /// Send the request and return a stream of SSE lines parsed into typed events.
105    ///
106    /// Uses manual SSE parsing because the Codex API does not return a
107    /// `Content-Type: text/event-stream` header, which `reqwest_eventsource`
108    /// (used by `async-openai`'s `create_stream`) requires.
109    async fn send_request(
110        &self,
111        request: Value,
112        headers: HeaderMap,
113    ) -> Result<impl futures::Stream<Item = Result<CodexStreamEvent>>> {
114        let url = format!("{}/responses", self.base_url);
115
116        debug!("Sending request to Codex API: {url}");
117        debug!(
118            "Codex request body: {}",
119            serde_json::to_string(&request).unwrap_or_else(|_| "<failed to serialize>".to_string())
120        );
121
122        let response = self.client.post(&url).headers(headers).json(&request).send().await?;
123
124        if !response.status().is_success() {
125            let status = response.status();
126            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
127
128            if matches!(status.as_u16(), 401 | 403) {
129                self.token_manager.clear_cache().await;
130            }
131
132            let message = format!("Codex API request failed with status {status}: {error_text}");
133            return Err(match status.as_u16() {
134                429 => LlmError::RateLimited(message),
135                s if (500..600).contains(&s) => LlmError::ServerError { status: Some(s), message },
136                _ => LlmError::ApiError(message),
137            });
138        }
139
140        let event_stream = response.bytes_stream().eventsource().filter_map(|result| {
141            std::future::ready(match result {
142                Ok(event) if event.data == "[DONE]" => None,
143                Ok(event) => match serde_json::from_str::<CodexStreamEvent>(&event.data) {
144                    Ok(parsed) => Some(Ok(parsed)),
145                    Err(e) => {
146                        debug!("Failed to parse Codex SSE line: {} - Error: {e}", event.data);
147                        None
148                    }
149                },
150                Err(e) => Some(Err(LlmError::StreamInterrupted(e.to_string()))),
151            })
152        });
153
154        Ok(event_stream)
155    }
156}
157
158impl StreamingModelProvider for CodexProvider {
159    fn model(&self) -> Option<crate::LlmModel> {
160        format!("{}:{}", super::PROVIDER_ID, self.model).parse().ok()
161    }
162
163    fn context_window(&self) -> Option<u32> {
164        get_context_window(super::PROVIDER_ID, &self.model)
165    }
166
167    fn stream_response(&self, context: &Context) -> LlmResponseStream {
168        let provider = self.clone();
169        let context = match self.model() {
170            Some(model) => context.filter_encrypted_reasoning(&model),
171            None => context.clone(),
172        };
173
174        Box::pin(async_stream::stream! {
175            let headers = match provider.build_headers().await {
176                Ok(h) => h,
177                Err(e) => {
178                    yield Err(e);
179                    return;
180                }
181            };
182
183            let request = match provider.build_wire_request(&context) {
184                Ok(r) => r,
185                Err(e) => {
186                    yield Err(e);
187                    return;
188                }
189            };
190
191            let event_stream = match provider.send_request(request, headers).await {
192                Ok(s) => s,
193                Err(e) => {
194                    yield Err(e);
195                    return;
196                }
197            };
198
199            let mut response_stream = Box::pin(process_response_stream(event_stream));
200            while let Some(result) = response_stream.next().await {
201                yield result;
202            }
203        })
204    }
205
206    fn display_name(&self) -> String {
207        format!("Codex ({})", self.model)
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214    use crate::ChatMessage;
215    use crate::ContentBlock;
216    use crate::ToolDefinition;
217    use crate::providers::test_capture_server::CaptureServer;
218    use crate::types::IsoString;
219    use aether_auth::{FakeOAuthCredentialStore, OAuthCredential};
220    use base64::Engine;
221    use base64::engine::general_purpose::URL_SAFE_NO_PAD;
222
223    #[test]
224    fn context_window_uses_codex_subscription_limit() {
225        let provider = create_test_provider();
226        assert_eq!(provider.context_window(), Some(272_000));
227    }
228
229    #[test]
230    fn display_name_includes_model() {
231        let provider = create_test_provider();
232        assert_eq!(provider.display_name(), "Codex (gpt-5.5)");
233    }
234
235    #[tokio::test]
236    async fn stream_response_sends_supported_protocol_version_for_gpt_5_6_luna() {
237        let mut server = CaptureServer::start().await;
238        let provider = server_backed_provider(&server).with_model("gpt-5.6-luna");
239        let mut context = Context::new(
240            vec![
241                ChatMessage::System { content: "You are helpful".to_string(), timestamp: IsoString::now() },
242                ChatMessage::User { content: vec![ContentBlock::text("Think harder")], timestamp: IsoString::now() },
243            ],
244            vec![ToolDefinition::new(
245                "bash",
246                "Run a command",
247                serde_json::from_str(r#"{"type": "object", "properties": {"cmd": {"type": "string"}}}"#).unwrap(),
248            )],
249        );
250        context.set_reasoning_effort(Some(crate::ReasoningEffort::Max));
251        context.set_prompt_cache_key(Some("session-abc".to_string()));
252
253        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
254        let captured = server.captured().await;
255
256        assert!(responses.iter().all(Result::is_ok), "{responses:?}");
257        assert_eq!(captured.body["reasoning"]["effort"], "max");
258        assert!(captured.body["reasoning"].get("context").is_none());
259        assert_eq!(captured.body["model"], "gpt-5.6-luna");
260        assert_eq!(captured.body["instructions"], "You are helpful");
261        assert_eq!(captured.body["tools"].as_array().unwrap().len(), 1);
262        assert!(captured.body.get("parallel_tool_calls").is_none());
263        assert_eq!(captured.body["input"][0]["role"], "user");
264        assert_eq!(captured.body["prompt_cache_key"], "session-abc");
265        assert_eq!(captured.body["store"], false);
266        assert_eq!(captured.body["stream"], true);
267        assert_eq!(captured.headers["chatgpt-account-id"], "account-1");
268        assert_eq!(captured.headers["version"], "0.144.0");
269        assert_eq!(captured.headers["accept"], "text/event-stream");
270        assert!(captured.headers.get("x-openai-internal-codex-responses-lite").is_none());
271        assert!(captured.headers.get("OpenAI-Beta").is_none());
272    }
273
274    #[tokio::test]
275    async fn stream_response_defaults_to_medium_effort_on_the_wire() {
276        let mut server = CaptureServer::start().await;
277        let provider = server_backed_provider(&server);
278        let context = Context::new(
279            vec![ChatMessage::User { content: vec![ContentBlock::text("Hello")], timestamp: IsoString::now() }],
280            vec![],
281        );
282
283        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
284        let captured = server.captured().await;
285
286        assert!(responses.iter().all(Result::is_ok), "{responses:?}");
287        assert_eq!(captured.body["reasoning"]["effort"], "medium");
288    }
289
290    fn server_backed_provider(server: &CaptureServer) -> CodexProvider {
291        let credential = OAuthCredential {
292            client_id: "test".to_string(),
293            access_token: test_jwt("account-1"),
294            refresh_token: None,
295            expires_at: Some(u64::MAX),
296            granted_scopes: Vec::new(),
297        };
298        let store: Arc<dyn OAuthCredentialStorage> =
299            Arc::new(FakeOAuthCredentialStore::new().with_credential("codex", credential));
300        CodexProvider::new(store).with_connection(crate::ProviderConnectionConfig {
301            base_url: Some(server.base_url.clone()),
302            ..Default::default()
303        })
304    }
305
306    fn test_jwt(account_id: &str) -> String {
307        let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none"}"#);
308        let payload = URL_SAFE_NO_PAD
309            .encode(serde_json::json!({"https://api.openai.com/auth": {"chatgpt_account_id": account_id}}).to_string());
310        format!("{header}.{payload}.signature")
311    }
312
313    fn create_test_provider() -> CodexProvider {
314        let store: Arc<dyn OAuthCredentialStorage> = Arc::new(FakeOAuthCredentialStore::new());
315        CodexProvider::new(store).with_model("gpt-5.5")
316    }
317}
llm/providers/codex/provider.rs

llm/providers/codex/
provider.rs
