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

llm/providers/codex/
provider.rs
