aither_codex/
client.rs

1//! [`LanguageModel`] implementation backed by the `ChatGPT` Codex Responses
2//! backend.
3//!
4//! [`Codex`] wraps [`aither_openai::Client`] pointed at
5//! `https://chatgpt.com/backend-api/codex` with subscription-auth request
6//! headers. It owns the OAuth token lifecycle: expired access tokens are
7//! refreshed lazily inside `respond`, a `401` mid-stream triggers one forced
8//! refresh and retry, and rotated credentials are pushed to a persistence
9//! callback so the refresh token (single-use) is never lost.
10
11use std::fmt;
12use std::sync::Arc;
13use std::time::Duration;
14
15use aither_core::llm::model::{Ability, Profile as ModelProfile};
16use aither_core::llm::{Event, LLMRequest, LanguageModel};
17use aither_openai::OpenAI;
18use async_lock::Mutex;
19use futures_lite::StreamExt as _;
20
21use crate::auth::{CodexCredentials, refresh_tokens};
22use crate::constant::{ACCOUNT_ID_HEADER, CODEX_BASE_URL, ORIGINATOR, RESIDENCY_HEADER};
23use crate::error::CodexError;
24
25/// Refresh the access token this long before its recorded expiry.
26const EXPIRY_SKEW: Duration = Duration::from_secs(60);
27
28/// Residency sentinel that means "no constraint" — the header is omitted.
29const NO_RESIDENCY_CONSTRAINT: &str = "no_constraint";
30
31/// Callback invoked whenever the credential set changes (token rotation).
32///
33/// The implementation is expected to persist the credentials; it is invoked
34/// synchronously while the token lock is held, so it should schedule any I/O
35/// rather than perform it inline.
36pub type OnRefresh = Arc<dyn Fn(&CodexCredentials) + Send + Sync>;
37
38/// A [`LanguageModel`] for a single model on the Codex subscription backend.
39#[derive(Clone)]
40pub struct Codex {
41    inner: Arc<Inner>,
42}
43
44struct Inner {
45    model: String,
46    base_url: String,
47    originator: String,
48    session_id: String,
49    cell: Mutex<Cell>,
50    on_refresh: Option<OnRefresh>,
51}
52
53struct Cell {
54    credentials: CodexCredentials,
55}
56
57impl Codex {
58    /// Create a client for `model` using the given credentials.
59    #[must_use]
60    pub fn new(credentials: CodexCredentials, model: impl Into<String>) -> Self {
61        Self {
62            inner: Arc::new(Inner {
63                model: model.into(),
64                base_url: CODEX_BASE_URL.into(),
65                originator: ORIGINATOR.into(),
66                session_id: session_id(),
67                cell: Mutex::new(Cell { credentials }),
68                on_refresh: None,
69            }),
70        }
71    }
72
73    /// Override the backend base URL (defaults to the Codex subscription
74    /// endpoint).
75    ///
76    /// # Panics
77    ///
78    /// Panics if the client has already been cloned/shared.
79    #[must_use]
80    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
81        Arc::get_mut(&mut self.inner)
82            .expect("with_base_url must be called before the client is shared")
83            .base_url = base_url.into();
84        self
85    }
86
87    /// Override the `originator` header (defaults to [`ORIGINATOR`]).
88    ///
89    /// # Panics
90    ///
91    /// Panics if the client has already been cloned/shared.
92    #[must_use]
93    pub fn with_originator(mut self, originator: impl Into<String>) -> Self {
94        Arc::get_mut(&mut self.inner)
95            .expect("with_originator must be called before the client is shared")
96            .originator = originator.into();
97        self
98    }
99
100    /// Register a callback fired whenever tokens rotate.
101    ///
102    /// # Panics
103    ///
104    /// Panics if the client has already been cloned/shared.
105    #[must_use]
106    pub fn on_refresh(
107        mut self,
108        callback: impl Fn(&CodexCredentials) + Send + Sync + 'static,
109    ) -> Self {
110        Arc::get_mut(&mut self.inner)
111            .expect("on_refresh must be called before the client is shared")
112            .on_refresh = Some(Arc::new(callback));
113        self
114    }
115
116    /// A snapshot of the current credentials.
117    pub async fn credentials(&self) -> CodexCredentials {
118        self.inner.cell.lock().await.credentials.clone()
119    }
120}
121
122impl Inner {
123    /// Build the inner `OpenAI` Responses client for a credential set.
124    fn openai_client(&self, credentials: &CodexCredentials) -> OpenAI {
125        let mut builder = OpenAI::builder(&credentials.access_token)
126            .base_url(&self.base_url)
127            .model(&self.model)
128            .store(false)
129            .extra_header("originator", &self.originator)
130            .extra_header("session-id", &self.session_id);
131        if let Some(account_id) = &credentials.account_id {
132            builder = builder.extra_header(ACCOUNT_ID_HEADER, account_id);
133        }
134        if let Some(residency) = credentials
135            .residency
136            .as_deref()
137            .filter(|r| *r != NO_RESIDENCY_CONSTRAINT)
138        {
139            builder = builder.extra_header(RESIDENCY_HEADER, residency);
140        }
141        builder.build()
142    }
143
144    /// Return valid credentials, refreshing if expired (or if `force`).
145    ///
146    /// The lock is held across the refresh so concurrent callers share one
147    /// refresh — the first to acquire performs it and the rest observe the
148    /// fresh credentials.
149    async fn valid_credentials(&self, force: bool) -> Result<CodexCredentials, CodexError> {
150        let mut cell = self.cell.lock().await;
151        if !force && !cell.credentials.is_expired(EXPIRY_SKEW) {
152            return Ok(cell.credentials.clone());
153        }
154        let fresh = refresh_tokens(&cell.credentials.refresh_token).await?;
155        tracing::debug!("refreshed ChatGPT access token");
156        if let Some(callback) = &self.on_refresh {
157            callback(&fresh);
158        }
159        cell.credentials = fresh.clone();
160        drop(cell);
161        Ok(fresh)
162    }
163
164    /// Build a request-ready client (fresh token + credential headers).
165    async fn request_client(&self, force_refresh: bool) -> Result<OpenAI, CodexError> {
166        let credentials = self.valid_credentials(force_refresh).await?;
167        Ok(self.openai_client(&credentials))
168    }
169}
170
171impl fmt::Debug for Codex {
172    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
173        f.debug_struct("Codex")
174            .field("model", &self.inner.model)
175            .field("base_url", &self.inner.base_url)
176            .field("session_id", &self.inner.session_id)
177            .finish_non_exhaustive()
178    }
179}
180
181impl LanguageModel for Codex {
182    type Error = CodexError;
183
184    fn respond(
185        &self,
186        request: LLMRequest,
187    ) -> impl futures_core::Stream<Item = Result<Event, Self::Error>> + Send {
188        let inner = Arc::clone(&self.inner);
189        async_stream::stream! {
190            let mut retried = false;
191            loop {
192                let client = match inner.request_client(retried).await {
193                    Ok(client) => client,
194                    Err(err) => {
195                        yield Err(err);
196                        return;
197                    }
198                };
199                let mut stream = std::pin::pin!(client.respond(request.clone()));
200                let mut unauthorized = false;
201                while let Some(event) = stream.next().await {
202                    match event {
203                        Err(err) if is_unauthorized(&err) && !retried => {
204                            unauthorized = true;
205                            break;
206                        }
207                        event => yield event.map_err(CodexError::OpenAI),
208                    }
209                }
210                if !unauthorized {
211                    return;
212                }
213                tracing::warn!("Codex request unauthorized; refreshing token and retrying");
214                retried = true;
215            }
216        }
217    }
218
219    fn profile(&self) -> impl core::future::Future<Output = ModelProfile> + Send {
220        let model = self.inner.model.clone();
221        async move {
222            let entry = aither_models::lookup(&model);
223            let mut profile = ModelProfile::new(
224                model.clone(),
225                "openai",
226                model,
227                "OpenAI model through a ChatGPT subscription",
228                entry
229                    .and_then(aither_models::ModelEntry::max_input_tokens)
230                    .unwrap_or(0),
231            )
232            .with_ability(Ability::ToolUse);
233            if let Some(entry) = entry {
234                for ability in entry.abilities() {
235                    if !profile.abilities.contains(ability) {
236                        profile.abilities.push(*ability);
237                    }
238                }
239            }
240            profile
241        }
242    }
243}
244
245/// Whether an `OpenAIError` is an HTTP 401.
246const fn is_unauthorized(err: &aither_openai::OpenAIError) -> bool {
247    matches!(
248        err,
249        aither_openai::OpenAIError::Http(zenwave::Error::Http { status, .. })
250            if status.as_u16() == 401
251    )
252}
253
254/// Random session identifier for the `session-id` header.
255fn session_id() -> String {
256    let mut bytes = [0u8; 16];
257    getrandom::fill(&mut bytes).expect("secure random source unavailable");
258    bytes.iter().fold(String::with_capacity(32), |mut s, b| {
259        use std::fmt::Write as _;
260        let _ = write!(s, "{b:02x}");
261        s
262    })
263}
aither_codex/client.rs

aither_codex/
client.rs
