Skip to main content

llm/providers/codex/
provider.rs

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