返回 CodeWhale
codex_model_cache.rs
根目录 / crates / tui / src / codex_model_cache.rs
1 //! Account-scoped roster for official Sign in with ChatGPT plan use.
2 //!
3 //! This replaces external Codex CLI cache/app-server discovery. Network access
4 //! uses the existing Codewhale provider client; only secret-free model metadata
5 //! is cached, keyed by the verified issuer, issued client ID, and subject.
6 //! A missing account roster offers no models. Public catalog rows and legacy
7 //! Codex credentials never prove permission to use a ChatGPT plan.
8
9 use std::io::Read;
10 use std::path::{Path, PathBuf};
11 use std::sync::Mutex;
12 use std::time::SystemTime;
13
14 #[cfg(unix)]
15 use std::os::unix::fs::OpenOptionsExt;
16
17 use chrono::{DateTime, Duration, Utc};
18 use serde::{Deserialize, Serialize};
19 use sha2::{Digest, Sha256};
20
21 use crate::config::{Config, ProviderKind};
22
23 const MAX_MODEL_CACHE_BYTES: u64 = 4 * 1024 * 1024;
24 const MODEL_CACHE_MAX_AGE: Duration = Duration::hours(24);
25 const MAX_FUTURE_CLOCK_SKEW: Duration = Duration::minutes(5);
26
27 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
28 pub(crate) enum CodexModelCacheFreshness {
29 Fresh,
30 Missing,
31 Stale,
32 Invalid,
33 }
34
35 impl CodexModelCacheFreshness {
36 #[must_use]
37 pub(crate) const fn picker_label(self) -> &'static str {
38 match self {
39 Self::Fresh => "ChatGPT OAuth",
40 Self::Missing => "OAuth roster missing",
41 Self::Stale => "OAuth roster stale",
42 Self::Invalid => "OAuth roster invalid",
43 }
44 }
45 }
46
47 #[derive(Debug, Clone, PartialEq, Eq)]
48 pub(crate) struct CodexModelRoster {
49 pub(crate) models: Vec<CodexModelMetadata>,
50 pub(crate) freshness: CodexModelCacheFreshness,
51 pub(crate) fetched_at: Option<DateTime<Utc>>,
52 pub(crate) observed_at: Option<DateTime<Utc>>,
53 pub(crate) source: &'static str,
54 pub(crate) observation_persisted: bool,
55 }
56
57 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
58 pub(crate) struct CodexModelMetadata {
59 pub(crate) id: String,
60 #[serde(default)]
61 pub(crate) display_name: Option<String>,
62 pub(crate) context_window: Option<u32>,
63 pub(crate) reasoning: Option<bool>,
64 pub(crate) efforts: Vec<String>,
65 }
66
67 impl CodexModelRoster {
68 fn fallback(freshness: CodexModelCacheFreshness, fetched_at: Option<DateTime<Utc>>) -> Self {
69 Self {
70 models: Vec::new(),
71 freshness,
72 fetched_at,
73 observed_at: None,
74 source: "chatgpt_plan_api",
75 observation_persisted: false,
76 }
77 }
78
79 #[must_use]
80 pub(crate) fn model_ids(&self) -> Vec<String> {
81 self.models.iter().map(|model| model.id.clone()).collect()
82 }
83
84 #[must_use]
85 pub(crate) fn metadata_for(&self, id: &str) -> Option<&CodexModelMetadata> {
86 self.models
87 .iter()
88 .find(|model| model.id.eq_ignore_ascii_case(id.trim()))
89 }
90
91 #[must_use]
92 pub(crate) fn preferred_model_id(&self) -> Option<&str> {
93 (self.freshness == CodexModelCacheFreshness::Fresh)
94 .then(|| self.models.first().map(|model| model.id.as_str()))
95 .flatten()
96 }
97 }
98
99 #[derive(Serialize, Deserialize)]
100 struct CatalogSnapshot {
101 fetched_at: DateTime<Utc>,
102 models: Vec<CodexModelMetadata>,
103 }
104
105 type RosterCacheKey = (PathBuf, Option<SystemTime>, u64);
106 static ROSTER_MEMO: Mutex<Option<(RosterCacheKey, CodexModelRoster)>> = Mutex::new(None);
107
108 /// An unscoped completion/catalog cannot borrow another account's roster.
109 #[must_use]
110 pub(crate) fn model_roster() -> CodexModelRoster {
111 CodexModelRoster::fallback(CodexModelCacheFreshness::Missing, None)
112 }
113
114 #[must_use]
115 pub(crate) fn model_roster_for(config: &Config) -> CodexModelRoster {
116 let Some(path) = snapshot_path(config) else {
117 return model_roster();
118 };
119 let key = match std::fs::symlink_metadata(&path) {
120 Ok(metadata) if !metadata.file_type().is_file() => {
121 return CodexModelRoster::fallback(CodexModelCacheFreshness::Invalid, None);
122 }
123 Ok(metadata) => (path.clone(), metadata.modified().ok(), metadata.len()),
124 Err(_) => (path.clone(), None, 0),
125 };
126 let now = Utc::now();
127 if let Ok(memo) = ROSTER_MEMO.lock()
128 && let Some((cached_key, roster)) = memo.as_ref()
129 && *cached_key == key
130 && roster.freshness == CodexModelCacheFreshness::Fresh
131 && roster
132 .fetched_at
133 .is_some_and(|fetched| now.signed_duration_since(fetched) <= MODEL_CACHE_MAX_AGE)
134 {
135 return roster.clone();
136 }
137 let roster = load_snapshot(&path, now);
138 if let Ok(mut memo) = ROSTER_MEMO.lock() {
139 *memo = Some((key, roster.clone()));
140 }
141 roster
142 }
143
144 fn registration_key(issuer: &str, client_id: &str, subject: &str) -> String {
145 let mut identity = Sha256::new();
146 identity.update(b"codewhale-chatgpt-plan-roster-v1\0");
147 for value in [issuer, client_id, subject] {
148 identity.update((value.len() as u64).to_le_bytes());
149 identity.update(value.as_bytes());
150 }
151 identity
152 .finalize()
153 .iter()
154 .map(|byte| format!("{byte:02x}"))
155 .collect()
156 }
157
158 fn snapshot_path(config: &Config) -> Option<PathBuf> {
159 let identity = config
160 .builtin_provider_identity(ProviderKind::OpenaiCodex)
161 .ok()?;
162 if config.provider_uses_custom_endpoint(&identity) {
163 return None;
164 }
165 let registration = crate::oauth::official_chatgpt_registration(config).ok()?;
166 let catalog_path = crate::models_dev_live::cache_path()?;
167 Some(catalog_path.parent()?.join(format!(
168 "chatgpt-plan-{}.json",
169 registration_key(
170 &registration.issuer,
171 &registration.client_id,
172 &registration.subject
173 )
174 )))
175 }
176
177 fn load_snapshot(path: &Path, now: DateTime<Utc>) -> CodexModelRoster {
178 let bytes = match read_cache_bytes(path) {
179 Ok(bytes) => bytes,
180 Err(freshness) => return CodexModelRoster::fallback(freshness, None),
181 };
182 let snapshot: CatalogSnapshot = match serde_json::from_slice(&bytes) {
183 Ok(snapshot) => snapshot,
184 Err(_) => return CodexModelRoster::fallback(CodexModelCacheFreshness::Invalid, None),
185 };
186 let age = now.signed_duration_since(snapshot.fetched_at);
187 if age < -MAX_FUTURE_CLOCK_SKEW
188 || snapshot.models.iter().any(|model| {
189 !crate::provider_lake::valid_catalog_model_id(&model.id)
190 || model
191 .display_name
192 .as_ref()
193 .is_some_and(|name| name.len() > 512 || name.chars().any(char::is_control))
194 || model.efforts.len() > 16
195 || model.efforts.iter().any(|effort| !valid_effort(effort))
196 || model
197 .context_window
198 .is_some_and(|window| !(1..=16_000_000).contains(&window))
199 })
200 {
201 return CodexModelRoster::fallback(
202 CodexModelCacheFreshness::Invalid,
203 Some(snapshot.fetched_at),
204 );
205 }
206 if age > MODEL_CACHE_MAX_AGE {
207 return CodexModelRoster::fallback(
208 CodexModelCacheFreshness::Stale,
209 Some(snapshot.fetched_at),
210 );
211 }
212 CodexModelRoster {
213 models: snapshot.models,
214 freshness: CodexModelCacheFreshness::Fresh,
215 fetched_at: Some(snapshot.fetched_at),
216 observed_at: None,
217 source: "chatgpt_plan_api",
218 observation_persisted: true,
219 }
220 }
221
222 fn valid_effort(effort: &str) -> bool {
223 !effort.is_empty()
224 && effort.len() <= 32
225 && effort
226 .bytes()
227 .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
228 }
229
230 /// Fetch through the existing provider client. A registration change during
231 /// the request cannot publish an old account's models under a new account.
232 pub(crate) async fn update_from_chatgpt(config: &Config) -> Result<CodexModelRoster, &'static str> {
233 let config = config.clone();
234 let prepared = config.clone();
235 #[cfg(test)]
236 let ticket = crate::test_support::env_scope_ticket();
237 let (path, client) = tokio::task::spawn_blocking(move || {
238 #[cfg(test)]
239 let _membership = crate::test_support::join_env_scope(ticket);
240 let path = snapshot_path(&prepared).ok_or("chatgpt_plan_permission_required")?;
241 let client = crate::client::CodewhaleClient::for_catalog_refresh(&prepared)
242 .map_err(|_| "chatgpt_plan_credentials_unavailable")?;
243 Ok::<_, &'static str>((path, client))
244 })
245 .await
246 .map_err(|_| "chatgpt_plan_credentials_unavailable")??;
247 let available = tokio::time::timeout(std::time::Duration::from_secs(20), client.list_models())
248 .await
249 .map_err(|_| "chatgpt_models_timeout")?
250 .map_err(|_| "chatgpt_models_unavailable")?;
251 // Catalog ordering and labels are provider facts. The basic official
252 // listing does not establish context limits or reasoning effort tiers.
253 if available.iter().any(|model| {
254 !crate::provider_lake::valid_catalog_model_id(&model.id)
255 || model
256 .display_name
257 .as_ref()
258 .is_some_and(|name| name.len() > 512 || name.chars().any(char::is_control))
259 }) {
260 return Err("chatgpt_models_invalid_response");
261 }
262 let models = available
263 .into_iter()
264 .map(|model| CodexModelMetadata {
265 id: model.id,
266 display_name: model.display_name,
267 context_window: None,
268 reasoning: None,
269 efforts: Vec::new(),
270 })
271 .collect();
272 let snapshot = CatalogSnapshot {
273 fetched_at: Utc::now(),
274 models,
275 };
276 #[cfg(test)]
277 let ticket = crate::test_support::env_scope_ticket();
278 tokio::task::spawn_blocking(move || {
279 #[cfg(test)]
280 let _membership = crate::test_support::join_env_scope(ticket);
281 if snapshot_path(&config).as_ref() != Some(&path) {
282 return Err("refresh_credentials_changed");
283 }
284 let encoded = serde_json::to_vec(&snapshot).map_err(|_| "cache_write_failed")?;
285 if encoded.len() as u64 > MAX_MODEL_CACHE_BYTES {
286 return Err("chatgpt_models_response_too_large");
287 }
288 codewhale_config::persistence::atomic_write(&path, &encoded)
289 .map_err(|_| "cache_write_failed")?;
290 if let Ok(mut memo) = ROSTER_MEMO.lock() {
291 *memo = None;
292 }
293 Ok(load_snapshot(&path, Utc::now()))
294 })
295 .await
296 .map_err(|_| "cache_write_failed")?
297 }
298
299 fn read_cache_bytes(path: &Path) -> Result<Vec<u8>, CodexModelCacheFreshness> {
300 let path_metadata = match std::fs::symlink_metadata(path) {
301 Ok(metadata) => metadata,
302 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
303 return Err(CodexModelCacheFreshness::Missing);
304 }
305 Err(_) => return Err(CodexModelCacheFreshness::Invalid),
306 };
307 if !path_metadata.file_type().is_file() || path_metadata.len() > MAX_MODEL_CACHE_BYTES {
308 return Err(CodexModelCacheFreshness::Invalid);
309 }
310 let mut file = match open_cache_file(path) {
311 Ok(file) => file,
312 Err(_) => return Err(CodexModelCacheFreshness::Invalid),
313 };
314 let metadata = match file.metadata() {
315 Ok(metadata) => metadata,
316 Err(_) => return Err(CodexModelCacheFreshness::Invalid),
317 };
318 if !metadata.file_type().is_file() || metadata.len() > MAX_MODEL_CACHE_BYTES {
319 return Err(CodexModelCacheFreshness::Invalid);
320 }
321
322 let mut bytes = Vec::with_capacity(metadata.len().min(MAX_MODEL_CACHE_BYTES) as usize);
323 if file
324 .by_ref()
325 .take(MAX_MODEL_CACHE_BYTES + 1)
326 .read_to_end(&mut bytes)
327 .is_err()
328 || bytes.len() as u64 > MAX_MODEL_CACHE_BYTES
329 {
330 return Err(CodexModelCacheFreshness::Invalid);
331 }
332 Ok(bytes)
333 }
334
335 #[cfg(test)]
336 pub(crate) fn install_test_chatgpt_roster(config: &Config, ids: &[&str]) -> anyhow::Result<()> {
337 install_test_chatgpt_roster_with_metadata(
338 config,
339 ids.iter()
340 .map(|id| CodexModelMetadata {
341 id: (*id).to_string(),
342 display_name: None,
343 context_window: None,
344 reasoning: None,
345 efforts: Vec::new(),
346 })
347 .collect(),
348 )
349 }
350
351 #[cfg(test)]
352 pub(crate) fn install_test_chatgpt_roster_with_metadata(
353 config: &Config,
354 models: Vec<CodexModelMetadata>,
355 ) -> anyhow::Result<()> {
356 let path = snapshot_path(config)
357 .ok_or_else(|| anyhow::anyhow!("test needs an owned ChatGPT registration"))?;
358 let snapshot = CatalogSnapshot {
359 fetched_at: Utc::now(),
360 models,
361 };
362 codewhale_config::persistence::atomic_write(&path, &serde_json::to_vec(&snapshot)?)?;
363 if let Ok(mut memo) = ROSTER_MEMO.lock() {
364 *memo = None;
365 }
366 Ok(())
367 }
368
369 fn open_cache_file(path: &Path) -> std::io::Result<std::fs::File> {
370 let mut options = std::fs::OpenOptions::new();
371 options.read(true);
372 #[cfg(unix)]
373 options.custom_flags(libc::O_NOFOLLOW);
374 options.open(path)
375 }
376
377 #[cfg(test)]
378 mod tests {
379 use super::*;
380
381 fn model(id: &str) -> CodexModelMetadata {
382 CodexModelMetadata {
383 id: id.to_string(),
384 display_name: Some(format!("Label for {id}")),
385 context_window: None,
386 reasoning: None,
387 efforts: Vec::new(),
388 }
389 }
390
391 fn save(path: &Path, fetched_at: DateTime<Utc>, models: Vec<CodexModelMetadata>) {
392 std::fs::write(
393 path,
394 serde_json::to_vec(&CatalogSnapshot { fetched_at, models }).unwrap(),
395 )
396 .unwrap();
397 }
398
399 #[test]
400 fn registration_scope_separates_accounts_workspaces_and_issuers() {
401 let key = registration_key("https://auth.openai.com", "oaiapp_workspace_a", "account_a");
402 for other in [
403 registration_key("https://auth.openai.com", "oaiapp_workspace_a", "account_b"),
404 registration_key("https://auth.openai.com", "oaiapp_workspace_b", "account_a"),
405 registration_key("https://other.example", "oaiapp_workspace_a", "account_a"),
406 ] {
407 assert_ne!(key, other);
408 }
409 assert!(!key.contains("account_a"));
410 assert!(!key.contains("workspace_a"));
411 }
412
413 #[test]
414 fn own_roster_follows_selected_registration_and_disappears_after_sign_out() {
415 let _env = crate::test_support::lock_test_env();
416 let directory = tempfile::tempdir().unwrap();
417 let directory_path = directory.path().canonicalize().unwrap();
418 let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", &directory_path);
419 let mut account_a = Config::default();
420 crate::oauth::install_test_chatgpt_registration_for(
421 &mut account_a,
422 "account-a",
423 "oaiapp_workspace_a",
424 )
425 .unwrap();
426 install_test_chatgpt_roster(&account_a, &["z-first", "a-second"]).unwrap();
427 let path_a = snapshot_path(&account_a).unwrap();
428 assert_eq!(
429 model_roster_for(&account_a).model_ids(),
430 ["z-first", "a-second"]
431 );
432
433 let mut account_b = Config::default();
434 crate::oauth::install_test_chatgpt_registration_for(
435 &mut account_b,
436 "account-b",
437 "oaiapp_workspace_a",
438 )
439 .unwrap();
440 assert_ne!(snapshot_path(&account_b).unwrap(), path_a);
441 assert_eq!(
442 model_roster_for(&account_b).freshness,
443 CodexModelCacheFreshness::Missing
444 );
445 install_test_chatgpt_roster(&account_b, &["b-only"]).unwrap();
446 assert_eq!(model_roster_for(&account_b).model_ids(), ["b-only"]);
447 assert_eq!(
448 model_roster_for(&account_a).model_ids(),
449 ["z-first", "a-second"]
450 );
451
452 let mut workspace_b = Config::default();
453 crate::oauth::install_test_chatgpt_registration_for(
454 &mut workspace_b,
455 "account-a",
456 "oaiapp_workspace_b",
457 )
458 .unwrap();
459 assert!(model_roster_for(&workspace_b).models.is_empty());
460 let generation = account_a
461 .provider_config_for(&account_a.test_identity_for_kind(ProviderKind::OpenaiCodex))
462 .unwrap()
463 .oauth_credential_generation
464 .as_ref()
465 .unwrap()
466 .clone();
467 let path = codewhale_config::chatgpt_oauth_generation_path(&generation).unwrap();
468 std::fs::remove_file(path).unwrap();
469 assert!(model_roster_for(&account_a).models.is_empty());
470 assert!(path_a.exists());
471 assert_eq!(model_roster_for(&account_b).model_ids(), ["b-only"]);
472 }
473
474 #[test]
475 fn own_snapshot_preserves_provider_order_and_labels_without_inventing_limits() {
476 let directory = tempfile::tempdir().unwrap();
477 let path = directory.path().join("roster.json");
478 let now = Utc::now();
479 save(&path, now, vec![model("z-first"), model("a-second")]);
480 let roster = load_snapshot(&path, now);
481 assert_eq!(roster.model_ids(), ["z-first", "a-second"]);
482 assert_eq!(roster.preferred_model_id(), Some("z-first"));
483 let metadata = roster.metadata_for("a-second").unwrap();
484 assert_eq!(metadata.display_name.as_deref(), Some("Label for a-second"));
485 assert_eq!(metadata.context_window, None);
486 assert_eq!(metadata.reasoning, None);
487 assert!(metadata.efforts.is_empty());
488 assert_eq!(roster.source, "chatgpt_plan_api");
489 assert!(roster.observation_persisted);
490 }
491
492 #[test]
493 fn missing_stale_future_and_unsafe_snapshots_offer_no_entitlements() {
494 let directory = tempfile::tempdir().unwrap();
495 let path = directory.path().join("roster.json");
496 let now = Utc::now();
497 assert!(load_snapshot(&path, now).model_ids().is_empty());
498 save(
499 &path,
500 now - MODEL_CACHE_MAX_AGE - Duration::seconds(1),
501 vec![model("old")],
502 );
503 let stale = load_snapshot(&path, now);
504 assert_eq!(stale.freshness, CodexModelCacheFreshness::Stale);
505 assert!(stale.models.is_empty());
506 save(
507 &path,
508 now + MAX_FUTURE_CLOCK_SKEW + Duration::seconds(1),
509 vec![model("future")],
510 );
511 assert_eq!(
512 load_snapshot(&path, now).freshness,
513 CodexModelCacheFreshness::Invalid
514 );
515 let mut unsafe_label = model("safe-id");
516 unsafe_label.display_name = Some("Unsafe\u{1b}[31m".to_string());
517 save(&path, now, vec![unsafe_label]);
518 assert_eq!(
519 load_snapshot(&path, now).freshness,
520 CodexModelCacheFreshness::Invalid
521 );
522 assert!(model_roster().models.is_empty());
523 }
524
525 #[cfg(unix)]
526 #[test]
527 fn cache_symlinks_are_rejected() {
528 let directory = tempfile::tempdir().unwrap();
529 let target = directory.path().join("target.json");
530 let path = directory.path().join("roster.json");
531 let now = Utc::now();
532 save(&target, now, vec![model("safe")]);
533 std::os::unix::fs::symlink(target, &path).unwrap();
534 assert_eq!(
535 load_snapshot(&path, now).freshness,
536 CodexModelCacheFreshness::Invalid
537 );
538 }
539 }
540
540 lines RUST