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