llm/providers/codex/
provider.rs

1use super::oauth::CodexTokenManager;
2use crate::catalog::Provider;
3use crate::provider::{LlmResponseStream, StreamingModelProvider, display_name, validate_reasoning};
4use crate::provider_connection::DEFAULT_STREAM_IDLE_TIMEOUT;
5use crate::providers::http::{SseData, http_client, post_sse, responses_code};
6use crate::providers::openai::responses::mappers::{ResponsesRequestPolicy, build_request};
7use crate::providers::openai::responses::streaming::decode_responses;
8use crate::providers::openai::responses::types::ResponsesRequest;
9use crate::providers::response_stream::{OpenedStream, response_stream};
10use crate::{Context, LlmModel, Result};
11use aether_auth::OAuthCredentialStorage;
12use reqwest::header::{AUTHORIZATION, HeaderMap, HeaderValue};
13use std::sync::Arc;
14use std::time::Duration;
15use tracing::debug;
16
17const CODEX_API_BASE: &str = "https://chatgpt.com/backend-api/codex";
18const CODEX_CLIENT_VERSION: &str = "0.159.0";
19
20#[derive(Clone)]
21pub struct CodexProvider {
22    base_url: String,
23    client: reqwest::Client,
24    model: String,
25    request_model: Option<String>,
26    token_manager: Arc<CodexTokenManager>,
27    idle_timeout: Duration,
28}
29
30impl CodexProvider {
31    pub fn new(store: Arc<dyn OAuthCredentialStorage>) -> Self {
32        Self::from_token_manager(CodexTokenManager::new(store, super::PROVIDER_ID))
33    }
34
35    pub(crate) fn from_token_manager(token_manager: CodexTokenManager) -> Self {
36        Self {
37            base_url: CODEX_API_BASE.to_string(),
38            client: http_client(),
39            model: "gpt-5.5".to_string(),
40            request_model: None,
41            token_manager: Arc::new(token_manager),
42            idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT,
43        }
44    }
45
46    pub fn with_connection(mut self, connection: crate::ProviderConnectionConfig) -> Self {
47        if let Some(base_url) = connection.base_url {
48            self.base_url = base_url.trim_end_matches('/').to_string();
49        }
50        self.request_model = connection.request_model;
51        self.idle_timeout = connection.idle_timeout;
52        self
53    }
54
55    pub fn with_model(mut self, model: &str) -> Self {
56        self.model = model.to_string();
57        self
58    }
59
60    pub async fn forget_cached_token(&self) {
61        self.token_manager.clear_cache().await;
62    }
63
64    fn build_wire_request(&self, context: &Context) -> Result<ResponsesRequest> {
65        let identity = self.model();
66        validate_reasoning(context, identity.as_ref())?;
67        let wire_model = self.request_model.as_deref().unwrap_or(&self.model);
68        build_request(identity.as_ref(), wire_model, context, &ResponsesRequestPolicy::CODEX)
69    }
70
71    async fn build_headers(&self) -> Result<HeaderMap> {
72        let (access_token, account_id) = self.token_manager.get_valid_token().await?;
73
74        let mut authorization = HeaderValue::from_str(&format!("Bearer {access_token}"))?;
75        authorization.set_sensitive(true);
76
77        let mut headers = HeaderMap::new();
78        headers.insert(AUTHORIZATION, authorization);
79        headers.insert("chatgpt-account-id", HeaderValue::from_str(&account_id)?);
80        headers.insert("originator", HeaderValue::from_static("codex_cli_rs"));
81        headers.insert("version", HeaderValue::from_static(CODEX_CLIENT_VERSION));
82
83        Ok(headers)
84    }
85
86    /// Send the request and return a stream of SSE lines parsed into typed events.
87    async fn send_request(&self, request: ResponsesRequest, headers: HeaderMap) -> Result<OpenedStream<SseData>> {
88        let url = format!("{}/responses", self.base_url);
89
90        debug!("Sending request to Codex API: {url}");
91        debug!(?request, "Codex request body");
92
93        match post_sse(&self.client, &url, headers, &request, responses_code).await {
94            Ok(opened) => Ok(opened),
95            Err(error) => {
96                if error.provider().map(|provider| provider.kind) == Some(crate::ProviderErrorKind::Authentication) {
97                    self.token_manager.clear_cache().await;
98                }
99                Err(error)
100            }
101        }
102    }
103}
104
105impl StreamingModelProvider for CodexProvider {
106    fn display_name(&self) -> String {
107        display_name(Provider::Codex, &self.model)
108    }
109
110    fn model(&self) -> Option<LlmModel> {
111        LlmModel::parse_for(Provider::Codex, &self.model).ok()
112    }
113
114    fn stream_response(&self, context: &Context) -> LlmResponseStream {
115        let provider = self.clone();
116        let context = context.clone();
117
118        response_stream(
119            async move {
120                let request = provider.build_wire_request(&context)?;
121                let headers = provider.build_headers().await?;
122                provider.send_request(request, headers).await
123            },
124            decode_responses(),
125            self.idle_timeout,
126        )
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use crate::ChatMessage;
134    use crate::LlmError;
135    use crate::ToolDefinition;
136    use crate::providers::test_capture_server::{CaptureServer, hello_context};
137    use aether_auth::{MemoryCredentialStore, 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(MemoryCredentialStore::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        assert_eq!(create_test_provider().with_model("gpt-custom").display_name(), "Codex (gpt-custom)");
169    }
170
171    #[tokio::test]
172    async fn stream_response_sends_supported_protocol_version_for_gpt_6_luna() {
173        let mut server = CaptureServer::start_responses().await;
174        let provider = server_backed_provider(&server).with_model("gpt-6-luna");
175        let mut context = Context::new(
176            vec![ChatMessage::system("You are helpful"), ChatMessage::user("Think harder")],
177            vec![ToolDefinition::new(
178                "bash",
179                "Run a command",
180                serde_json::from_str(r#"{"type": "object", "properties": {"cmd": {"type": "string"}}}"#).unwrap(),
181            )],
182        );
183        context.set_reasoning_effort(crate::ReasoningEffort::Max);
184        context.set_prompt_cache_key(Some("session-abc".to_string()));
185
186        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
187        let captured = server.captured().await;
188
189        assert!(responses.iter().all(Result::is_ok), "{responses:?}");
190        assert_eq!(captured.body["reasoning"]["effort"], "max");
191        assert!(captured.body["reasoning"].get("context").is_none());
192        assert_eq!(captured.body["model"], "gpt-6-luna");
193        assert_eq!(captured.body["instructions"], "You are helpful");
194        assert_eq!(captured.body["tools"].as_array().unwrap().len(), 1);
195        assert!(captured.body["tools"][0].get("strict").is_none());
196        assert_eq!(captured.body["parallel_tool_calls"], true);
197        assert_eq!(captured.body["input"][0]["role"], "user");
198        assert_eq!(captured.body["prompt_cache_key"], "session-abc");
199        assert_eq!(captured.body["store"], false);
200        assert_eq!(captured.body["stream"], true);
201        assert_eq!(captured.headers["chatgpt-account-id"], "account-1");
202        assert_eq!(captured.headers["version"], "0.159.0");
203        assert_eq!(captured.headers["accept"], "text/event-stream");
204        assert!(captured.headers.get("x-openai-internal-codex-responses-lite").is_none());
205        assert!(captured.headers.get("OpenAI-Beta").is_none());
206    }
207
208    #[tokio::test]
209    async fn stream_response_defaults_to_medium_effort_on_the_wire() {
210        let mut server = CaptureServer::start_responses().await;
211        let provider = server_backed_provider(&server);
212        let context = Context::new(vec![ChatMessage::user("Hello")], vec![]);
213
214        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
215        let captured = server.captured().await;
216
217        assert!(responses.iter().all(Result::is_ok), "{responses:?}");
218        assert_eq!(captured.body["reasoning"]["effort"], "medium");
219        assert_eq!(captured.body["reasoning"]["summary"], "auto");
220    }
221
222    #[tokio::test]
223    async fn request_model_routes_the_request_without_changing_catalog_identity() {
224        let mut server = CaptureServer::start_responses().await;
225        let provider = authenticated_provider().with_connection(crate::ProviderConnectionConfig {
226            base_url: Some(server.base_url.clone()),
227            request_model: Some("codex-deployment".to_string()),
228            ..Default::default()
229        });
230
231        let responses = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
232        let captured = server.captured().await;
233
234        assert!(responses.iter().all(Result::is_ok), "{responses:?}");
235        assert_eq!(captured.body["model"], "codex-deployment");
236        assert_eq!(provider.model().unwrap().to_string(), "codex:gpt-5.5");
237    }
238
239    #[tokio::test]
240    async fn a_logout_reaches_a_provider_that_forgets_its_cached_token() {
241        let mut server = CaptureServer::start_responses().await;
242        let store = authenticated_store();
243        let provider = CodexProvider::new(Arc::clone(&store)).with_connection(crate::ProviderConnectionConfig {
244            base_url: Some(server.base_url.clone()),
245            ..Default::default()
246        });
247
248        let before = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
249        server.captured().await;
250        store.delete("codex").await.unwrap();
251        provider.forget_cached_token().await;
252        let after = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
253
254        assert!(before.iter().all(Result::is_ok), "{before:?}");
255        assert!(matches!(after.as_slice(), [Err(LlmError::OAuthError(_))]), "{after:?}");
256    }
257
258    fn server_backed_provider(server: &CaptureServer) -> CodexProvider {
259        authenticated_provider().with_connection(crate::ProviderConnectionConfig {
260            base_url: Some(server.base_url.clone()),
261            ..Default::default()
262        })
263    }
264
265    fn authenticated_provider() -> CodexProvider {
266        CodexProvider::new(authenticated_store())
267    }
268
269    fn authenticated_store() -> Arc<dyn OAuthCredentialStorage> {
270        let credential = OAuthCredential {
271            client_id: "test".to_string(),
272            access_token: test_jwt("account-1"),
273            refresh_token: None,
274            expires_at: Some(u64::MAX),
275        };
276        Arc::new(MemoryCredentialStore::new().with_credential("codex", credential))
277    }
278
279    fn test_jwt(account_id: &str) -> String {
280        let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none"}"#);
281        let payload = URL_SAFE_NO_PAD
282            .encode(serde_json::json!({"https://api.openai.com/auth": {"chatgpt_account_id": account_id}}).to_string());
283        format!("{header}.{payload}.signature")
284    }
285
286    fn create_test_provider() -> CodexProvider {
287        let store: Arc<dyn OAuthCredentialStorage> = Arc::new(MemoryCredentialStore::new());
288        CodexProvider::new(store).with_model("gpt-5.5")
289    }
290}
llm/providers/codex/provider.rs

llm/providers/codex/
provider.rs
