atman_runtime/providers/
codex.rs

1use serde::{Deserialize, Serialize};
2use tokio::sync::broadcast;
3use tokio_util::sync::CancellationToken;
4
5use crate::error::RuntimeError;
6use crate::event::{NodeEvent, Observable, TurnId};
7use crate::message::{ImageData, Message, MessagePart, MessageRole};
8use crate::provider::{
9    AssistantMessage, CallTiming, DEFAULT_STREAM_BUFFER, LlmRequest, Provider, StopReason,
10    TokenUsage, estimate_tokens,
11};
12use crate::tool::BoxFut;
13use anyhow::Context;
14
15const CODEX_BASE: &str = "https://chatgpt.com/backend-api/codex";
16
17/// ChatGPT backend provider. Requires `originator: codex_cli_rs` header for Cloudflare.
18pub struct CodexProvider {
19    name: String,
20    access_token: String,
21    account_id: String,
22    client: reqwest::Client,
23}
24
25impl CodexProvider {
26    pub fn new(
27        name: impl Into<String>,
28        access_token: impl Into<String>,
29        account_id: impl Into<String>,
30    ) -> Self {
31        Self {
32            name: name.into(),
33            access_token: access_token.into(),
34            account_id: account_id.into(),
35            client: reqwest::Client::new(),
36        }
37    }
38
39    fn build_body(&self, req: &LlmRequest) -> ResponsesRequest {
40        let model = req
41            .model
42            .split_once('/')
43            .map(|(_, slug)| slug)
44            .or_else(|| req.model.split_once(':').map(|(_, slug)| slug))
45            .unwrap_or(&req.model)
46            .to_string();
47
48        let input = build_input_items(req);
49        let tools = build_tools(&req.tools);
50
51        ResponsesRequest {
52            model,
53            input,
54            instructions: req.system.clone(),
55            tools,
56            stream: true,
57            store: false,
58            reasoning: Some(ReasoningConfig {
59                effort: Some("medium".into()),
60                summary: "auto".into(),
61            }),
62            text: Some(TextConfig {
63                verbosity: "medium".into(),
64            }),
65            include: Some(vec!["reasoning.encrypted_content".into()]),
66        }
67    }
68
69    fn build_request(&self, req: &LlmRequest) -> reqwest::RequestBuilder {
70        let body = self.build_body(req);
71        self.client
72            .post(format!("{CODEX_BASE}/responses"))
73            .bearer_auth(&self.access_token)
74            .header("chatgpt-account-id", &self.account_id)
75            .header("originator", "codex_cli_rs")
76            .header("OpenAI-Beta", "responses=experimental")
77            .header("accept", "text/event-stream")
78            .json(&body)
79    }
80}
81
82fn build_input_items(req: &LlmRequest) -> Vec<InputItem> {
83    let mut tool_names: std::collections::HashMap<String, String> =
84        std::collections::HashMap::new();
85    for m in &req.messages {
86        if m.role == MessageRole::Assistant {
87            for p in &m.parts {
88                if let MessagePart::ToolUse { id, name, .. } = p {
89                    // Codex Responses API requires "_" in function names.
90                    tool_names.insert(id.clone(), name.replace('.', "_"));
91                }
92            }
93        }
94    }
95
96    let mut items: Vec<InputItem> = Vec::new();
97
98    for m in &req.messages {
99        match m.role {
100            MessageRole::User => {
101                let content = build_user_content(&m.parts);
102                items.push(InputItem {
103                    role: Some("user".into()),
104                    content: Some(content),
105                    item_type: Some("message".into()),
106                    call_id: None,
107                    name: None,
108                    arguments: None,
109                    output: None,
110                });
111            }
112            MessageRole::Assistant => {
113                let (text, tool_calls) = split_assistant_parts(&m.parts);
114                if let Some(t) = text {
115                    items.push(InputItem {
116                        role: Some("assistant".into()),
117                        content: Some(t),
118                        item_type: Some("message".into()),
119                        call_id: None,
120                        name: None,
121                        arguments: None,
122                        output: None,
123                    });
124                }
125                for tc in tool_calls {
126                    items.push(InputItem {
127                        role: None,
128                        content: None,
129                        item_type: Some("function_call".into()),
130                        call_id: Some(tc.id),
131                        name: Some(tc.name),
132                        arguments: Some(tc.arguments),
133                        output: None,
134                    });
135                }
136            }
137            MessageRole::Tool => {
138                for p in &m.parts {
139                    if let MessagePart::ToolResult {
140                        tool_use_id,
141                        content,
142                        ..
143                    } = p
144                    {
145                        let name = tool_names.get(tool_use_id).cloned();
146                        items.push(InputItem {
147                            role: None,
148                            content: None,
149                            item_type: Some("function_call_output".into()),
150                            call_id: Some(tool_use_id.clone()),
151                            name,
152                            arguments: None,
153                            output: Some(content.clone()),
154                        });
155                    }
156                }
157            }
158            MessageRole::System => {}
159        }
160    }
161
162    items
163}
164
165fn build_user_content(parts: &[MessagePart]) -> String {
166    let mut parts_out: Vec<serde_json::Value> = Vec::new();
167    for p in parts {
168        match p {
169            MessagePart::Text { text } => {
170                parts_out.push(serde_json::json!({"type": "input_text", "text": text}));
171            }
172            MessagePart::Image { source } => {
173                let data = match &source.data {
174                    ImageData::Base64 { data } => data.clone(),
175                    ImageData::Path { path } => {
176                        let bytes = std::fs::read(path).unwrap_or_default();
177                        use base64::Engine;
178                        base64::engine::general_purpose::STANDARD.encode(&bytes)
179                    }
180                };
181                parts_out.push(serde_json::json!({
182                    "type": "input_image",
183                    "image_url": format!("data:{};base64,{}", source.media_type, data)
184                }));
185            }
186            MessagePart::CompactSummary { summary, .. } => {
187                parts_out.push(serde_json::json!({"type": "input_text", "text": summary}));
188            }
189            _ => {}
190        }
191    }
192    if parts_out.len() == 1
193        && parts_out[0].get("type").and_then(|v| v.as_str()) == Some("input_text")
194    {
195        parts_out[0]["text"].as_str().unwrap_or("").to_string()
196    } else {
197        serde_json::to_string(&parts_out).unwrap_or_default()
198    }
199}
200
201struct AssistantSplit {
202    id: String,
203    name: String,
204    arguments: String,
205}
206
207fn split_assistant_parts(parts: &[MessagePart]) -> (Option<String>, Vec<AssistantSplit>) {
208    let mut text = String::new();
209    let mut tools: Vec<AssistantSplit> = Vec::new();
210    for p in parts {
211        match p {
212            MessagePart::Text { text: t } => text.push_str(t),
213            MessagePart::ToolUse { id, name, input } => tools.push(AssistantSplit {
214                id: id.clone(),
215                // Codex Responses API requires "_" in function names.
216                name: name.replace('.', "_"),
217                arguments: serde_json::to_string(input).unwrap_or_default(),
218            }),
219            _ => {}
220        }
221    }
222    let text_out = if text.is_empty() { None } else { Some(text) };
223    (text_out, tools)
224}
225
226fn build_tools(tools: &[crate::tool::ToolSpec]) -> Vec<ResponsesTool> {
227    tools
228        .iter()
229        .map(|t| ResponsesTool {
230            r#type: "function".into(),
231            // Codex Responses API rejects "." in function names (e.g. "fs.read").
232            // Replace with "_" for outbound, revert on inbound.
233            name: t.name.replace('.', "_"),
234            description: t.description.clone(),
235            parameters: t.input_schema.clone(),
236        })
237        .collect()
238}
239
240impl Provider for CodexProvider {
241    fn name(&self) -> &str {
242        &self.name
243    }
244
245    fn call<'a>(&'a self, req: LlmRequest) -> BoxFut<'a, Result<AssistantMessage, RuntimeError>> {
246        // The Codex backend always uses streaming (Responses API with store=false).
247        // We call streaming internally and collect the result.
248        let observable = self.call_streaming(req);
249        Box::pin(async move {
250            // Drop the broadcast events — callers of `call()` don't consume them.
251            let _events = observable.events;
252            observable.output.await
253        })
254    }
255
256    fn call_streaming(&self, req: LlmRequest) -> Observable<AssistantMessage> {
257        let request = self.build_request(&req);
258        let turn_id = turn_id_from_req(&req);
259        let (tx, events) = broadcast::channel(DEFAULT_STREAM_BUFFER);
260        let cancel = CancellationToken::new();
261        let cancel_for_task = cancel.clone();
262
263        let output: BoxFut<'static, Result<AssistantMessage, RuntimeError>> =
264            Box::pin(async move {
265                use eventsource_stream::Eventsource;
266                use futures::StreamExt;
267
268                let resp = tokio::select! {
269                    biased;
270                    _ = cancel_for_task.cancelled() => {
271                        return Err(RuntimeError::Cancelled("codex cancelled before send".into()));
272                    }
273                    r = request.send() => r.map_err(net_err)?,
274                };
275                let status = resp.status();
276                if !status.is_success() {
277                    let body_text = resp.text().await.unwrap_or_default();
278                    return Err(RuntimeError::ToolFailed(format!(
279                        "codex http {status}: {body_text}"
280                    )));
281                }
282
283                let mut stream = resp.bytes_stream().eventsource();
284
285                let mut acc_text = String::new();
286                let mut acc_thinking = String::new();
287                let mut cumulative = 0u64;
288                let mut final_usage: Option<ResponsesUsage> = None;
289                let mut resp_model: Option<String> = None;
290                let mut resp_id: Option<String> = None;
291                let mut stop_reason = StopReason::End;
292
293                let mut partial_tool_calls: Vec<PartialToolCall> = Vec::new();
294
295                while let Some(event) = tokio::select! {
296                    biased;
297                    _ = cancel_for_task.cancelled() => None,
298                    next = stream.next() => next,
299                } {
300                    let event =
301                        event.map_err(|e| RuntimeError::ToolFailed(format!("codex sse: {e}")))?;
302                    if event.data.is_empty() || event.data == "[DONE]" {
303                        continue;
304                    }
305                    let parsed: serde_json::Value = match serde_json::from_str(&event.data) {
306                        Ok(v) => v,
307                        Err(_) => continue,
308                    };
309
310                    let ev_type = parsed["type"].as_str().unwrap_or("");
311
312                    match ev_type {
313                        "response.output_text.delta" => {
314                            if let Some(delta) = parsed["delta"].as_str() {
315                                acc_text.push_str(delta);
316                                cumulative += estimate_tokens(delta);
317                                let _ = tx.send(NodeEvent::LlmChunk {
318                                    text: delta.to_string(),
319                                    cumulative_tokens: cumulative,
320                                });
321                            }
322                        }
323
324                        "response.reasoning_text.delta" => {
325                            if let Some(delta) = parsed["delta"].as_str() {
326                                acc_thinking.push_str(delta);
327                                let _ = tx.send(NodeEvent::ThinkingChunk {
328                                    text: delta.to_string(),
329                                });
330                            }
331                        }
332
333                        "response.output_item.added" => {
334                            if let Some(item) = parsed.get("item")
335                                && item["type"].as_str() == Some("function_call")
336                            {
337                                let idx = parsed["output_index"].as_u64().unwrap_or(0) as usize;
338                                while partial_tool_calls.len() <= idx {
339                                    partial_tool_calls.push(PartialToolCall::default());
340                                }
341                                let slot = &mut partial_tool_calls[idx];
342                                slot.id = item["call_id"].as_str().unwrap_or("").to_string();
343                                // Convert "_" back to "." for atman tool names.
344                                slot.name = item["name"].as_str().unwrap_or("").replace('_', ".");
345                            }
346                        }
347
348                        "response.function_call_arguments.delta" => {
349                            let idx = parsed["output_index"].as_u64().unwrap_or(0) as usize;
350                            while partial_tool_calls.len() <= idx {
351                                partial_tool_calls.push(PartialToolCall::default());
352                            }
353                            if let Some(delta) = parsed["delta"].as_str() {
354                                partial_tool_calls[idx].arguments.push_str(delta);
355                            }
356                        }
357
358                        "response.completed" => {
359                            if let Some(r) = parsed.get("response") {
360                                resp_model = r["model"].as_str().map(|s| s.to_string());
361                                resp_id = r["id"].as_str().map(|s| s.to_string());
362                                if let Some(u) = r.get("usage") {
363                                    final_usage =
364                                        serde_json::from_value::<ResponsesUsage>(u.clone()).ok();
365                                }
366                                if r["status"].as_str() == Some("cancelled") {
367                                    stop_reason = StopReason::Cancelled;
368                                }
369                            }
370                        }
371
372                        "error" => {
373                            let msg = parsed["message"].as_str().unwrap_or("unknown codex error");
374                            return Err(RuntimeError::ToolFailed(msg.to_string()));
375                        }
376
377                        _ => {}
378                    }
379                }
380
381                if cancel_for_task.is_cancelled() {
382                    let _ = tx.send(NodeEvent::LlmDone {
383                        total_tokens: cumulative,
384                    });
385                    return Err(RuntimeError::Cancelled("codex cancelled mid-stream".into()));
386                }
387
388                let total_output = final_usage
389                    .as_ref()
390                    .and_then(|u| u.output_tokens)
391                    .unwrap_or(cumulative);
392                let _ = tx.send(NodeEvent::LlmDone {
393                    total_tokens: total_output,
394                });
395
396                let mut parts: Vec<MessagePart> = Vec::new();
397                if !acc_thinking.is_empty() {
398                    parts.push(MessagePart::Thinking {
399                        thinking: acc_thinking,
400                        signature: None,
401                    });
402                }
403                if !acc_text.is_empty() {
404                    parts.push(MessagePart::Text { text: acc_text });
405                }
406                for tc in partial_tool_calls {
407                    if tc.name.is_empty() {
408                        continue;
409                    }
410                    let input: serde_json::Value = if tc.arguments.is_empty() {
411                        serde_json::Value::Object(Default::default())
412                    } else {
413                        serde_json::from_str(&tc.arguments).unwrap_or(serde_json::Value::Null)
414                    };
415                    parts.push(MessagePart::ToolUse {
416                        id: tc.id,
417                        name: tc.name,
418                        input,
419                    });
420                }
421
422                let token_usage = final_usage.map(|u| TokenUsage {
423                    input: u.input_tokens.unwrap_or(0),
424                    cached_input: u
425                        .input_tokens_details
426                        .as_ref()
427                        .and_then(|d| d.cached_tokens)
428                        .unwrap_or(0),
429                    output: u.output_tokens.unwrap_or(0),
430                    cache_write: u
431                        .input_tokens_details
432                        .as_ref()
433                        .and_then(|d| d.cache_write_tokens)
434                        .unwrap_or(0),
435                    reasoning_tokens: u
436                        .output_tokens_details
437                        .as_ref()
438                        .and_then(|d| d.reasoning_tokens)
439                        .unwrap_or(0),
440                });
441
442                Ok(AssistantMessage {
443                    message: Message {
444                        role: MessageRole::Assistant,
445                        parts,
446                        turn_id,
447                    },
448                    stop_reason,
449                    token_usage: token_usage.unwrap_or_default(),
450                    timing: CallTiming::default(),
451                    model: resp_model.unwrap_or_default(),
452                    response_id: resp_id,
453                })
454            });
455        Observable {
456            output,
457            events,
458            cancel,
459        }
460    }
461
462    fn discover_models(
463        &self,
464    ) -> crate::tool::BoxFut<'static, Vec<crate::provider::DiscoveredModel>> {
465        let access_token = self.access_token.clone();
466        let account_id = self.account_id.clone();
467        Box::pin(async move {
468            let client = reqwest::Client::new();
469            let resp = match client
470                .get("https://chatgpt.com/backend-api/wham/models")
471                .query(&[("client_version", "0.0.0")])
472                .bearer_auth(&access_token)
473                .header("ChatGPT-Account-Id", &account_id)
474                .send()
475                .await
476            {
477                Ok(r) => r,
478                Err(e) => {
479                    crate::notify!(
480                        warn,
481                        location = Inline,
482                        stack = dedupe("codex.models.fetch_failed", 60_000),
483                        "fetch codex models failed: {e:#}"
484                    );
485                    return vec![];
486                }
487            };
488            let Ok(body) = resp.json::<serde_json::Value>().await else {
489                return vec![];
490            };
491            let Some(list) = body["models"].as_array() else {
492                return vec![];
493            };
494            list.iter()
495                .filter_map(|m| {
496                    let slug = format!("codex/{}", m["slug"].as_str()?);
497                    if slug.is_empty() {
498                        return None;
499                    }
500                    let context_budget = m["context_window"].as_u64();
501                    let thinking = m["supported_reasoning_levels"]
502                        .as_array()
503                        .map(|a: &Vec<serde_json::Value>| !a.is_empty())
504                        .unwrap_or(false);
505                    Some(crate::provider::DiscoveredModel {
506                        slug,
507                        context_budget,
508                        thinking,
509                    })
510                })
511                .collect()
512        })
513    }
514}
515
516const CODEX_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
517const CODEX_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
518const CODEX_AUTHORIZE_URL: &str = "https://auth.openai.com/oauth/authorize";
519const CODEX_REDIRECT_URI: &str = "http://localhost:1455/auth/callback";
520
521impl crate::oauth::OAuthProvider for CodexProvider {
522    fn authorize_url() -> (String, crate::oauth::Pkce, String) {
523        let pkce = crate::oauth::Pkce::generate();
524        let state = crate::oauth::generate_state();
525        let url = format!(
526            "{}?response_type=code&client_id={}&redirect_uri={}&code_challenge={}&code_challenge_method=S256&state={}&scope=openid+profile+email+offline_access",
527            CODEX_AUTHORIZE_URL, CODEX_CLIENT_ID, CODEX_REDIRECT_URI, pkce.challenge, state
528        );
529        (url, pkce, state)
530    }
531
532    fn exchange_code(
533        code: &str,
534        verifier: &str,
535    ) -> std::pin::Pin<
536        Box<dyn std::future::Future<Output = anyhow::Result<crate::oauth::TokenResult>> + Send>,
537    > {
538        let code = code.to_string();
539        let verifier = verifier.to_string();
540        Box::pin(async move {
541            let client = reqwest::Client::new();
542            let resp = client
543                .post(CODEX_TOKEN_URL)
544                .form(&[
545                    ("grant_type", "authorization_code"),
546                    ("code", &code),
547                    ("redirect_uri", CODEX_REDIRECT_URI),
548                    ("client_id", CODEX_CLIENT_ID),
549                    ("code_verifier", &verifier),
550                ])
551                .send()
552                .await
553                .context("token exchange request")?;
554
555            let status = resp.status();
556            let body_text = resp.text().await.unwrap_or_default();
557            if !status.is_success() {
558                anyhow::bail!("token exchange failed (HTTP {status}): {body_text}");
559            }
560
561            #[derive(serde::Deserialize)]
562            struct R {
563                access_token: String,
564                refresh_token: Option<String>,
565                id_token: Option<String>,
566            }
567            let data: R = serde_json::from_str(&body_text).context("parse token response")?;
568
569            let expires_at = crate::oauth::parse_jwt_exp(&data.access_token)
570                .unwrap_or_else(|| chrono::Utc::now().timestamp() + 3600);
571            let account = data
572                .id_token
573                .as_deref()
574                .and_then(crate::oauth::extract_account_from_id_token);
575
576            Ok(crate::oauth::TokenResult {
577                access_token: data.access_token,
578                refresh_token: data.refresh_token,
579                expires_at,
580                account,
581            })
582        })
583    }
584
585    fn refresh_token(
586        token: &str,
587    ) -> std::pin::Pin<
588        Box<dyn std::future::Future<Output = anyhow::Result<crate::oauth::TokenResult>> + Send>,
589    > {
590        let token = token.to_string();
591        Box::pin(async move {
592            let client = reqwest::Client::new();
593            let resp = client
594                .post(CODEX_TOKEN_URL)
595                .form(&[
596                    ("grant_type", "refresh_token"),
597                    ("refresh_token", &token),
598                    ("client_id", CODEX_CLIENT_ID),
599                ])
600                .send()
601                .await
602                .context("token refresh request")?;
603
604            let status = resp.status();
605            let body_text = resp.text().await.unwrap_or_default();
606            if !status.is_success() {
607                anyhow::bail!("token refresh failed (HTTP {status}): {body_text}");
608            }
609
610            #[derive(serde::Deserialize)]
611            struct R {
612                access_token: String,
613                refresh_token: Option<String>,
614                id_token: Option<String>,
615            }
616            let data: R = serde_json::from_str(&body_text).context("parse refresh response")?;
617
618            let expires_at = crate::oauth::parse_jwt_exp(&data.access_token)
619                .unwrap_or_else(|| chrono::Utc::now().timestamp() + 3600);
620            let account = data
621                .id_token
622                .as_deref()
623                .and_then(crate::oauth::extract_account_from_id_token);
624
625            Ok(crate::oauth::TokenResult {
626                access_token: data.access_token,
627                refresh_token: data.refresh_token,
628                expires_at,
629                account,
630            })
631        })
632    }
633
634    fn from_stored(stored: &crate::auth_store::StoredProvider) -> Self {
635        let account_id = stored.account.as_deref().unwrap_or("");
636        CodexProvider::new(&stored.name, &stored.access_token, account_id)
637    }
638}
639
640#[derive(Default)]
641struct PartialToolCall {
642    id: String,
643    name: String,
644    arguments: String,
645}
646
647fn turn_id_from_req(req: &LlmRequest) -> TurnId {
648    req.messages
649        .first()
650        .map(|m| m.turn_id.clone())
651        .unwrap_or_else(TurnId::now)
652}
653
654fn net_err(e: reqwest::Error) -> RuntimeError {
655    RuntimeError::ToolFailed(format!("codex net: {e}"))
656}
657
658#[derive(Serialize)]
659struct ResponsesRequest {
660    model: String,
661    input: Vec<InputItem>,
662    #[serde(skip_serializing_if = "Option::is_none")]
663    instructions: Option<String>,
664    #[serde(skip_serializing_if = "Vec::is_empty")]
665    tools: Vec<ResponsesTool>,
666    stream: bool,
667    store: bool,
668    #[serde(skip_serializing_if = "Option::is_none")]
669    reasoning: Option<ReasoningConfig>,
670    #[serde(skip_serializing_if = "Option::is_none")]
671    text: Option<TextConfig>,
672    #[serde(skip_serializing_if = "Option::is_none")]
673    include: Option<Vec<String>>,
674}
675
676#[derive(Serialize)]
677struct InputItem {
678    #[serde(skip_serializing_if = "Option::is_none")]
679    role: Option<String>,
680    #[serde(skip_serializing_if = "Option::is_none")]
681    content: Option<String>,
682    #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
683    item_type: Option<String>,
684    #[serde(skip_serializing_if = "Option::is_none")]
685    call_id: Option<String>,
686    #[serde(skip_serializing_if = "Option::is_none")]
687    name: Option<String>,
688    #[serde(skip_serializing_if = "Option::is_none")]
689    arguments: Option<String>,
690    #[serde(skip_serializing_if = "Option::is_none")]
691    output: Option<String>,
692}
693
694#[derive(Serialize)]
695struct ResponsesTool {
696    #[serde(rename = "type")]
697    r#type: String,
698    name: String,
699    #[serde(skip_serializing_if = "Option::is_none")]
700    description: Option<String>,
701    parameters: serde_json::Value,
702}
703
704#[derive(Serialize)]
705struct ReasoningConfig {
706    #[serde(skip_serializing_if = "Option::is_none")]
707    effort: Option<String>,
708    summary: String,
709}
710
711#[derive(Serialize)]
712struct TextConfig {
713    verbosity: String,
714}
715
716#[derive(Deserialize, Default)]
717struct ResponsesUsage {
718    #[serde(default)]
719    input_tokens: Option<u64>,
720    #[serde(default)]
721    output_tokens: Option<u64>,
722    #[serde(default)]
723    input_tokens_details: Option<InputTokensDetails>,
724    #[serde(default)]
725    output_tokens_details: Option<OutputTokensDetails>,
726}
727
728#[derive(Deserialize, Default)]
729struct InputTokensDetails {
730    #[serde(default)]
731    cached_tokens: Option<u64>,
732    #[serde(default)]
733    cache_write_tokens: Option<u64>,
734}
735
736#[derive(Deserialize, Default)]
737struct OutputTokensDetails {
738    #[serde(default)]
739    reasoning_tokens: Option<u64>,
740}
atman_runtime/providers/codex.rs

atman_runtime/providers/
codex.rs
