mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
load models from disk and set a ttl and etag (#7722)
# External (non-OpenAI) Pull Request Requirements Before opening this Pull Request, please read the dedicated "Contributing" markdown file or your PR may be closed: https://github.com/openai/codex/blob/main/docs/contributing.md If your PR conforms to our contribution guidelines, replace this text with a detailed and high quality description of your changes. Include a link to a bug report or enhancement request.
This commit is contained in:
@@ -32,7 +32,9 @@ use crate::token_data::TokenData;
|
||||
use crate::token_data::parse_id_token;
|
||||
use crate::util::try_parse_error_message;
|
||||
use codex_protocol::account::PlanType as AccountPlanType;
|
||||
use once_cell::sync::Lazy;
|
||||
use serde_json::Value;
|
||||
use tempfile::TempDir;
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -62,6 +64,8 @@ const REFRESH_TOKEN_UNKNOWN_MESSAGE: &str =
|
||||
const REFRESH_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
|
||||
pub const REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR: &str = "CODEX_REFRESH_TOKEN_URL_OVERRIDE";
|
||||
|
||||
static TEST_AUTH_TEMP_DIRS: Lazy<Mutex<Vec<TempDir>>> = Lazy::new(|| Mutex::new(Vec::new()));
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum RefreshTokenError {
|
||||
#[error("{0}")]
|
||||
@@ -1088,11 +1092,19 @@ impl AuthManager {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
#[expect(clippy::expect_used)]
|
||||
/// Create an AuthManager with a specific CodexAuth, for testing only.
|
||||
pub fn from_auth_for_testing(auth: CodexAuth) -> Arc<Self> {
|
||||
let cached = CachedAuth { auth: Some(auth) };
|
||||
let temp_dir = tempfile::tempdir().expect("temp codex home");
|
||||
let codex_home = temp_dir.path().to_path_buf();
|
||||
TEST_AUTH_TEMP_DIRS
|
||||
.lock()
|
||||
.expect("lock test codex homes")
|
||||
.push(temp_dir);
|
||||
Arc::new(Self {
|
||||
codex_home: PathBuf::new(),
|
||||
codex_home,
|
||||
inner: RwLock::new(cached),
|
||||
enable_codex_api_key_env: false,
|
||||
auth_credentials_store_mode: AuthCredentialsStoreMode::File,
|
||||
@@ -1104,6 +1116,10 @@ impl AuthManager {
|
||||
self.inner.read().ok().and_then(|c| c.auth.clone())
|
||||
}
|
||||
|
||||
pub fn codex_home(&self) -> &Path {
|
||||
&self.codex_home
|
||||
}
|
||||
|
||||
/// Force a reload of the auth information from auth.json. Returns
|
||||
/// whether the auth value changed.
|
||||
pub fn reload(&self) -> bool {
|
||||
|
||||
@@ -51,6 +51,7 @@ impl ConversationManager {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
/// Construct with a dummy AuthManager containing the provided CodexAuth.
|
||||
/// Used for integration tests: should not be used by ordinary business logic.
|
||||
pub fn with_auth(auth: CodexAuth) -> Self {
|
||||
@@ -213,7 +214,7 @@ impl ConversationManager {
|
||||
}
|
||||
|
||||
pub async fn list_models(&self) -> Vec<ModelPreset> {
|
||||
self.models_manager.available_models.read().await.clone()
|
||||
self.models_manager.list_models().await
|
||||
}
|
||||
|
||||
pub fn get_models_manager(&self) -> Arc<ModelsManager> {
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
use chrono::DateTime;
|
||||
use chrono::Utc;
|
||||
use codex_protocol::openai_models::ModelInfo;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use std::io;
|
||||
use std::io::ErrorKind;
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
use tokio::fs;
|
||||
|
||||
/// Serialized snapshot of models and metadata cached on disk.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub(crate) struct ModelsCache {
|
||||
pub(crate) fetched_at: DateTime<Utc>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) etag: Option<String>,
|
||||
pub(crate) models: Vec<ModelInfo>,
|
||||
}
|
||||
|
||||
impl ModelsCache {
|
||||
/// Returns `true` when the cache entry has not exceeded the configured TTL.
|
||||
pub(crate) fn is_fresh(&self, ttl: Duration) -> bool {
|
||||
if ttl.is_zero() {
|
||||
return false;
|
||||
}
|
||||
let Ok(ttl_duration) = chrono::Duration::from_std(ttl) else {
|
||||
return false;
|
||||
};
|
||||
let age = Utc::now().signed_duration_since(self.fetched_at);
|
||||
age <= ttl_duration
|
||||
}
|
||||
}
|
||||
|
||||
/// Read and deserialize the cache file if it exists.
|
||||
pub(crate) async fn load_cache(path: &Path) -> io::Result<Option<ModelsCache>> {
|
||||
match fs::read(path).await {
|
||||
Ok(contents) => {
|
||||
let cache = serde_json::from_slice(&contents)
|
||||
.map_err(|err| io::Error::new(ErrorKind::InvalidData, err.to_string()))?;
|
||||
Ok(Some(cache))
|
||||
}
|
||||
Err(err) if err.kind() == ErrorKind::NotFound => Ok(None),
|
||||
Err(err) => Err(err),
|
||||
}
|
||||
}
|
||||
|
||||
/// Persist the cache contents to disk, creating parent directories as needed.
|
||||
pub(crate) async fn save_cache(path: &Path, cache: &ModelsCache) -> io::Result<()> {
|
||||
if let Some(parent) = path.parent() {
|
||||
fs::create_dir_all(parent).await?;
|
||||
}
|
||||
let json = serde_json::to_vec_pretty(cache)
|
||||
.map_err(|err| io::Error::new(ErrorKind::InvalidData, err.to_string()))?;
|
||||
fs::write(path, json).await
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
mod cache;
|
||||
pub mod model_family;
|
||||
pub mod model_presets;
|
||||
pub mod models_manager;
|
||||
|
||||
@@ -1,11 +1,19 @@
|
||||
use chrono::Utc;
|
||||
use codex_api::ModelsClient;
|
||||
use codex_api::ReqwestTransport;
|
||||
use codex_protocol::openai_models::ModelInfo;
|
||||
use codex_protocol::openai_models::ModelPreset;
|
||||
use codex_protocol::openai_models::ModelsResponse;
|
||||
use http::HeaderMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::RwLock;
|
||||
use tokio::sync::TryLockError;
|
||||
use tracing::error;
|
||||
|
||||
use super::cache;
|
||||
use super::cache::ModelsCache;
|
||||
use crate::api_bridge::auth_provider_from_auth;
|
||||
use crate::api_bridge::map_api_error;
|
||||
use crate::auth::AuthManager;
|
||||
@@ -17,29 +25,41 @@ use crate::openai_models::model_family::ModelFamily;
|
||||
use crate::openai_models::model_family::find_family_for_model;
|
||||
use crate::openai_models::model_presets::builtin_model_presets;
|
||||
|
||||
const MODEL_CACHE_FILE: &str = "models_cache.json";
|
||||
const DEFAULT_MODEL_CACHE_TTL: Duration = Duration::from_secs(300);
|
||||
|
||||
/// Coordinates remote model discovery plus cached metadata on disk.
|
||||
#[derive(Debug)]
|
||||
pub struct ModelsManager {
|
||||
// todo(aibrahim) merge available_models and model family creation into one struct
|
||||
pub available_models: RwLock<Vec<ModelPreset>>,
|
||||
pub remote_models: RwLock<Vec<ModelInfo>>,
|
||||
pub etag: String,
|
||||
pub auth_manager: Arc<AuthManager>,
|
||||
available_models: RwLock<Vec<ModelPreset>>,
|
||||
remote_models: RwLock<Vec<ModelInfo>>,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
etag: RwLock<Option<String>>,
|
||||
codex_home: PathBuf,
|
||||
cache_ttl: Duration,
|
||||
}
|
||||
|
||||
impl ModelsManager {
|
||||
/// Construct a manager scoped to the provided `AuthManager`.
|
||||
pub fn new(auth_manager: Arc<AuthManager>) -> Self {
|
||||
let codex_home = auth_manager.codex_home().to_path_buf();
|
||||
Self {
|
||||
available_models: RwLock::new(builtin_model_presets(auth_manager.get_auth_mode())),
|
||||
remote_models: RwLock::new(Vec::new()),
|
||||
etag: String::new(),
|
||||
auth_manager,
|
||||
etag: RwLock::new(None),
|
||||
codex_home,
|
||||
cache_ttl: DEFAULT_MODEL_CACHE_TTL,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn refresh_available_models(
|
||||
&self,
|
||||
provider: &ModelProviderInfo,
|
||||
) -> CoreResult<Vec<ModelInfo>> {
|
||||
/// Fetch the latest remote models, using the on-disk cache when still fresh.
|
||||
pub async fn refresh_available_models(&self, provider: &ModelProviderInfo) -> CoreResult<()> {
|
||||
if self.try_load_cache().await {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let auth = self.auth_manager.auth();
|
||||
let api_provider = provider.to_api_provider(auth.as_ref().map(|auth| auth.mode))?;
|
||||
let api_auth = auth_provider_from_auth(auth.clone(), provider).await?;
|
||||
@@ -50,21 +70,30 @@ impl ModelsManager {
|
||||
if client_version == "0.0.0" {
|
||||
client_version = "99.99.99";
|
||||
}
|
||||
let response = client
|
||||
let ModelsResponse { models, etag } = client
|
||||
.list_models(client_version, HeaderMap::new())
|
||||
.await
|
||||
.map_err(map_api_error)?;
|
||||
|
||||
let models = response.models;
|
||||
*self.remote_models.write().await = models.clone();
|
||||
let available_models = self.build_available_models().await;
|
||||
{
|
||||
let mut available_models_guard = self.available_models.write().await;
|
||||
*available_models_guard = available_models;
|
||||
}
|
||||
Ok(models)
|
||||
let etag = (!etag.is_empty()).then_some(etag);
|
||||
|
||||
self.apply_remote_models(models.clone()).await;
|
||||
*self.etag.write().await = etag.clone();
|
||||
self.persist_cache(&models, etag).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn list_models(&self) -> Vec<ModelPreset> {
|
||||
self.available_models.read().await.clone()
|
||||
}
|
||||
|
||||
pub fn try_list_models(&self) -> Result<Vec<ModelPreset>, TryLockError> {
|
||||
self.available_models
|
||||
.try_read()
|
||||
.map(|models| models.clone())
|
||||
}
|
||||
|
||||
/// Look up the requested model family while applying remote metadata overrides.
|
||||
pub async fn construct_model_family(&self, model: &str, config: &Config) -> ModelFamily {
|
||||
find_family_for_model(model)
|
||||
.with_config_overrides(config)
|
||||
@@ -72,11 +101,55 @@ impl ModelsManager {
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
/// Offline helper that builds a `ModelFamily` without consulting remote state.
|
||||
pub fn construct_model_family_offline(model: &str, config: &Config) -> ModelFamily {
|
||||
find_family_for_model(model).with_config_overrides(config)
|
||||
}
|
||||
|
||||
async fn build_available_models(&self) -> Vec<ModelPreset> {
|
||||
/// Replace the cached remote models and rebuild the derived presets list.
|
||||
async fn apply_remote_models(&self, models: Vec<ModelInfo>) {
|
||||
*self.remote_models.write().await = models;
|
||||
self.build_available_models().await;
|
||||
}
|
||||
|
||||
/// Attempt to satisfy the refresh from the cache when it matches the provider and TTL.
|
||||
async fn try_load_cache(&self) -> bool {
|
||||
let cache_path = self.cache_path();
|
||||
let cache = match cache::load_cache(&cache_path).await {
|
||||
Ok(cache) => cache,
|
||||
Err(err) => {
|
||||
error!("failed to load models cache: {err}");
|
||||
return false;
|
||||
}
|
||||
};
|
||||
let cache = match cache {
|
||||
Some(cache) => cache,
|
||||
None => return false,
|
||||
};
|
||||
if !cache.is_fresh(self.cache_ttl) {
|
||||
return false;
|
||||
}
|
||||
let models = cache.models.clone();
|
||||
*self.etag.write().await = cache.etag.clone();
|
||||
self.apply_remote_models(models.clone()).await;
|
||||
true
|
||||
}
|
||||
|
||||
/// Serialize the latest fetch to disk for reuse across future processes.
|
||||
async fn persist_cache(&self, models: &[ModelInfo], etag: Option<String>) {
|
||||
let cache = ModelsCache {
|
||||
fetched_at: Utc::now(),
|
||||
etag,
|
||||
models: models.to_vec(),
|
||||
};
|
||||
let cache_path = self.cache_path();
|
||||
if let Err(err) = cache::save_cache(&cache_path, &cache).await {
|
||||
error!("failed to write models cache: {err}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert remote model metadata into picker-ready presets, marking defaults.
|
||||
async fn build_available_models(&self) {
|
||||
let mut available_models = self.remote_models.read().await.clone();
|
||||
available_models.sort_by(|a, b| b.priority.cmp(&a.priority));
|
||||
let mut model_presets: Vec<ModelPreset> = available_models
|
||||
@@ -87,22 +160,29 @@ impl ModelsManager {
|
||||
if let Some(default) = model_presets.first_mut() {
|
||||
default.is_default = true;
|
||||
}
|
||||
model_presets
|
||||
{
|
||||
let mut available_models_guard = self.available_models.write().await;
|
||||
*available_models_guard = model_presets;
|
||||
}
|
||||
}
|
||||
|
||||
fn cache_path(&self) -> PathBuf {
|
||||
self.codex_home.join(MODEL_CACHE_FILE)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::cache::ModelsCache;
|
||||
use super::*;
|
||||
use crate::CodexAuth;
|
||||
use crate::auth::AuthCredentialsStoreMode;
|
||||
use crate::model_provider_info::WireApi;
|
||||
use codex_protocol::openai_models::ModelsResponse;
|
||||
use core_test_support::responses::mount_models_once;
|
||||
use serde_json::json;
|
||||
use wiremock::Mock;
|
||||
use tempfile::tempdir;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::method;
|
||||
use wiremock::matchers::path;
|
||||
|
||||
fn remote_model(slug: &str, display: &str, priority: i32) -> ModelInfo {
|
||||
serde_json::from_value(json!({
|
||||
@@ -146,35 +226,28 @@ mod tests {
|
||||
remote_model("priority-low", "Low", 1),
|
||||
remote_model("priority-high", "High", 10),
|
||||
];
|
||||
let response = ModelsResponse {
|
||||
models: remote_models.clone(),
|
||||
};
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/models"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "application/json")
|
||||
.set_body_json(&response),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let models_mock = mount_models_once(
|
||||
&server,
|
||||
ModelsResponse {
|
||||
models: remote_models.clone(),
|
||||
etag: String::new(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth_manager =
|
||||
AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
|
||||
let manager = ModelsManager::new(auth_manager);
|
||||
let provider = provider_for(server.uri());
|
||||
|
||||
let returned = manager
|
||||
manager
|
||||
.refresh_available_models(&provider)
|
||||
.await
|
||||
.expect("refresh succeeds");
|
||||
|
||||
assert_eq!(returned, remote_models);
|
||||
let cached_remote = manager.remote_models.read().await.clone();
|
||||
assert_eq!(cached_remote, remote_models);
|
||||
|
||||
let available = manager.available_models.read().await.clone();
|
||||
let available = manager.list_models().await;
|
||||
assert_eq!(available.len(), 2);
|
||||
assert_eq!(available[0].model, "priority-high");
|
||||
assert!(
|
||||
@@ -183,5 +256,128 @@ mod tests {
|
||||
);
|
||||
assert_eq!(available[1].model, "priority-low");
|
||||
assert!(!available[1].is_default);
|
||||
assert_eq!(
|
||||
models_mock.requests().len(),
|
||||
1,
|
||||
"expected a single /models request"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_available_models_uses_cache_when_fresh() {
|
||||
let server = MockServer::start().await;
|
||||
let remote_models = vec![remote_model("cached", "Cached", 5)];
|
||||
let models_mock = mount_models_once(
|
||||
&server,
|
||||
ModelsResponse {
|
||||
models: remote_models.clone(),
|
||||
etag: String::new(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let codex_home = tempdir().expect("temp dir");
|
||||
let auth_manager = Arc::new(AuthManager::new(
|
||||
codex_home.path().to_path_buf(),
|
||||
false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
));
|
||||
let manager = ModelsManager::new(auth_manager);
|
||||
let provider = provider_for(server.uri());
|
||||
|
||||
manager
|
||||
.refresh_available_models(&provider)
|
||||
.await
|
||||
.expect("first refresh succeeds");
|
||||
assert_eq!(
|
||||
*manager.remote_models.read().await,
|
||||
remote_models,
|
||||
"remote cache should store fetched models"
|
||||
);
|
||||
|
||||
// Second call should read from cache and avoid the network.
|
||||
manager
|
||||
.refresh_available_models(&provider)
|
||||
.await
|
||||
.expect("cached refresh succeeds");
|
||||
assert_eq!(
|
||||
*manager.remote_models.read().await,
|
||||
remote_models,
|
||||
"cache path should not mutate stored models"
|
||||
);
|
||||
assert_eq!(
|
||||
models_mock.requests().len(),
|
||||
1,
|
||||
"cache hit should avoid a second /models request"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_available_models_refetches_when_cache_stale() {
|
||||
let server = MockServer::start().await;
|
||||
let initial_models = vec![remote_model("stale", "Stale", 1)];
|
||||
let initial_mock = mount_models_once(
|
||||
&server,
|
||||
ModelsResponse {
|
||||
models: initial_models.clone(),
|
||||
etag: String::new(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let codex_home = tempdir().expect("temp dir");
|
||||
let auth_manager = Arc::new(AuthManager::new(
|
||||
codex_home.path().to_path_buf(),
|
||||
false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
));
|
||||
let manager = ModelsManager::new(auth_manager);
|
||||
let provider = provider_for(server.uri());
|
||||
|
||||
manager
|
||||
.refresh_available_models(&provider)
|
||||
.await
|
||||
.expect("initial refresh succeeds");
|
||||
|
||||
// Rewrite cache with an old timestamp so it is treated as stale.
|
||||
let cache_path = codex_home.path().join(MODEL_CACHE_FILE);
|
||||
let contents =
|
||||
std::fs::read_to_string(&cache_path).expect("cache file should exist after refresh");
|
||||
let mut cache: ModelsCache =
|
||||
serde_json::from_str(&contents).expect("cache should deserialize");
|
||||
cache.fetched_at = Utc::now() - chrono::Duration::hours(1);
|
||||
std::fs::write(&cache_path, serde_json::to_string_pretty(&cache).unwrap())
|
||||
.expect("cache rewrite succeeds");
|
||||
|
||||
let updated_models = vec![remote_model("fresh", "Fresh", 9)];
|
||||
server.reset().await;
|
||||
let refreshed_mock = mount_models_once(
|
||||
&server,
|
||||
ModelsResponse {
|
||||
models: updated_models.clone(),
|
||||
etag: String::new(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
manager
|
||||
.refresh_available_models(&provider)
|
||||
.await
|
||||
.expect("second refresh succeeds");
|
||||
assert_eq!(
|
||||
*manager.remote_models.read().await,
|
||||
updated_models,
|
||||
"stale cache should trigger refetch"
|
||||
);
|
||||
assert_eq!(
|
||||
initial_mock.requests().len(),
|
||||
1,
|
||||
"initial refresh should only hit /models once"
|
||||
);
|
||||
assert_eq!(
|
||||
refreshed_mock.requests().len(),
|
||||
1,
|
||||
"stale cache refresh should fetch /models once"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user