Skip to main content

imp_llm/providers/
openai_codex.rs

1use std::pin::Pin;
2
3use async_trait::async_trait;
4use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
5use futures_core::Stream;
6use serde_json::Value;
7
8use crate::auth::{ApiKey, AuthStore};
9use crate::error::{Error, Result};
10use crate::model::{Model, ModelMeta};
11use crate::provider::{Context, Provider, RequestOptions};
12use crate::stream::StreamEvent;
13
14use super::openai::{build_request_json, stream_response_json};
15
16const CODEX_API_URL: &str = "https://chatgpt.com/backend-api/codex/responses";
17const JWT_CLAIM_PATH: &str = "https://api.openai.com/auth";
18
19/// ChatGPT/Codex-backed OpenAI provider.
20pub struct OpenAiCodexProvider {
21    client: reqwest::Client,
22    models: Vec<ModelMeta>,
23}
24
25impl Default for OpenAiCodexProvider {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31impl OpenAiCodexProvider {
32    pub fn new() -> Self {
33        Self {
34            client: super::streaming_http_client(),
35            models: crate::model::builtin_openai_codex_models(),
36        }
37    }
38}
39
40fn extract_account_id(token: &str) -> Result<String> {
41    let payload = token
42        .split('.')
43        .nth(1)
44        .ok_or_else(|| Error::Auth("Invalid ChatGPT OAuth token".into()))?;
45    let decoded = URL_SAFE_NO_PAD
46        .decode(payload)
47        .map_err(|_| Error::Auth("Failed to decode ChatGPT OAuth token".into()))?;
48    let claims: Value = serde_json::from_slice(&decoded)
49        .map_err(|_| Error::Auth("Failed to parse ChatGPT OAuth token claims".into()))?;
50
51    claims
52        .get(JWT_CLAIM_PATH)
53        .and_then(|value| value.get("chatgpt_account_id"))
54        .and_then(Value::as_str)
55        .filter(|value| !value.is_empty())
56        .map(str::to_string)
57        .ok_or_else(|| Error::Auth("ChatGPT OAuth token is missing chatgpt_account_id".into()))
58}
59
60fn strip_unsupported_codex_fields(request: &mut Value) {
61    let Some(object) = request.as_object_mut() else {
62        return;
63    };
64
65    for field in [
66        "context_management",
67        "max_completion_tokens",
68        "max_output_tokens",
69        "max_tokens",
70        "metadata",
71        "temperature",
72        "user",
73    ] {
74        object.remove(field);
75    }
76}
77
78fn add_codex_request_fields(request: &mut Value, session_id: Option<&str>) {
79    strip_unsupported_codex_fields(request);
80
81    let Some(object) = request.as_object_mut() else {
82        return;
83    };
84
85    object.insert("store".into(), Value::Bool(false));
86    object.insert("parallel_tool_calls".into(), Value::Bool(true));
87    object.insert("tool_choice".into(), Value::String("auto".into()));
88    object.insert(
89        "include".into(),
90        Value::Array(vec![Value::String("reasoning.encrypted_content".into())]),
91    );
92    object.insert(
93        "text".into(),
94        serde_json::json!({
95            "verbosity": "medium",
96        }),
97    );
98    if let Some(session_id) = session_id.filter(|value| !value.is_empty()) {
99        object.insert(
100            "prompt_cache_key".into(),
101            Value::String(session_id.to_string()),
102        );
103    }
104}
105
106fn build_headers(
107    account_id: &str,
108    api_key: &str,
109    session_id: Option<&str>,
110) -> Vec<(String, String)> {
111    let mut headers = vec![
112        ("authorization".to_string(), format!("Bearer {api_key}")),
113        ("chatgpt-account-id".to_string(), account_id.to_string()),
114        ("originator".to_string(), "imp".to_string()),
115        (
116            "OpenAI-Beta".to_string(),
117            "responses=experimental".to_string(),
118        ),
119        ("accept".to_string(), "text/event-stream".to_string()),
120        ("content-type".to_string(), "application/json".to_string()),
121        (
122            "user-agent".to_string(),
123            format!("imp/{}", env!("CARGO_PKG_VERSION")),
124        ),
125    ];
126
127    if let Some(session_id) = session_id.filter(|value| !value.is_empty()) {
128        headers.push(("session_id".to_string(), session_id.to_string()));
129    }
130
131    headers
132}
133
134#[async_trait]
135impl Provider for OpenAiCodexProvider {
136    fn stream(
137        &self,
138        model: &Model,
139        context: Context,
140        options: RequestOptions,
141        api_key: &str,
142    ) -> Pin<Box<dyn Stream<Item = Result<StreamEvent>> + Send>> {
143        let account_id = match extract_account_id(api_key) {
144            Ok(account_id) => account_id,
145            Err(error) => {
146                return Box::pin(futures::stream::once(async move { Err(error) }));
147            }
148        };
149
150        let mut request = build_request_json(model, context, options);
151        add_codex_request_fields(&mut request, None);
152        let headers = build_headers(&account_id, api_key, None);
153        stream_response_json(
154            self.client.clone(),
155            CODEX_API_URL.to_string(),
156            headers,
157            request,
158        )
159    }
160
161    async fn resolve_auth(&self, auth: &AuthStore) -> Result<ApiKey> {
162        if let Some(oauth) = auth.get_oauth("openai") {
163            return Ok(oauth.access_token.clone());
164        }
165        if let Some(oauth) = auth.get_oauth("openai-codex") {
166            return Ok(oauth.access_token.clone());
167        }
168        Err(Error::Auth(
169            "No ChatGPT OAuth credential found. Run `imp login openai`.".into(),
170        ))
171    }
172
173    fn id(&self) -> &str {
174        "openai-codex"
175    }
176
177    fn models(&self) -> &[ModelMeta] {
178        &self.models
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[test]
187    fn codex_request_strips_unsupported_fields() {
188        let mut request = serde_json::json!({
189            "model": "gpt-5.4",
190            "max_output_tokens": 1234,
191            "temperature": 0.2,
192            "metadata": {"source": "test"},
193        });
194
195        add_codex_request_fields(&mut request, None);
196
197        let object = request.as_object().expect("request object");
198        assert!(!object.contains_key("max_output_tokens"));
199        assert!(!object.contains_key("temperature"));
200        assert!(!object.contains_key("metadata"));
201        assert_eq!(object.get("store"), Some(&Value::Bool(false)));
202        assert_eq!(
203            object.get("tool_choice"),
204            Some(&Value::String("auto".into()))
205        );
206    }
207
208    #[test]
209    fn codex_request_leaves_max_output_tokens_absent_when_unset() {
210        let mut request = serde_json::json!({
211            "model": "gpt-5.4"
212        });
213
214        add_codex_request_fields(&mut request, None);
215
216        let object = request.as_object().expect("request object");
217        assert!(!object.contains_key("max_output_tokens"));
218    }
219}