fix: use model snapshots for account test models

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-06-30 00:00:24 +08:00
co-authored by Codex
parent 4f6b873c67
commit d461cb6325
9 changed files with 95 additions and 101 deletions
+2 -5
View File
@@ -8,7 +8,7 @@ use axum::{
HeaderMap, StatusCode,
},
response::{IntoResponse, Response},
Extension, Json,
Json,
};
use std::collections::HashMap;
@@ -27,7 +27,6 @@ use crate::{
},
monitoring::{billing, service::AdminUsageRecord},
},
http::middleware::request_id::RequestId,
infra::{
json::{clamp_limit, clamp_page, total_pages, Page},
time::{china_datetime, china_relative_time, china_rfc3339},
@@ -706,16 +705,14 @@ pub(crate) async fn test_account_connection(
pub(crate) async fn account_models(
State(state): State<AppState>,
Query(query): Query<AccountIdQuery>,
Extension(request_id): Extension<RequestId>,
headers: HeaderMap,
) -> Result<impl IntoResponse, AdminError> {
let request_id = request_id.as_str().to_string();
let account_id = query.id;
require_admin_auth(&state, &headers).await?;
match state
.services
.admin_accounts
.account_models(&account_id, &request_id)
.account_models(&account_id)
.await
{
Ok(models) => Ok(AdminResponse::new(
+4
View File
@@ -20,6 +20,7 @@ use crate::{
store::SqliteAccountStore,
token_refresh::{RuntimeRefreshPolicy, TokenRefresher},
},
upstream::models::ModelService,
upstream::transport::CodexBackendClient,
};
@@ -34,6 +35,7 @@ pub struct AdminAccountService {
pub store: SqliteAccountStore,
pub(crate) cookies: SqliteCookieStore,
pub(crate) codex: StdArc<CodexBackendClient>,
pub(crate) models: StdArc<ModelService>,
pub(crate) account_pool: StdArc<RuntimeAccountPoolService>,
pub(crate) token_refresher: StdArc<dyn TokenRefresher>,
pub(crate) oauth: oauth::AccountOAuthService,
@@ -45,6 +47,7 @@ pub struct AdminAccountServiceParts {
pub store: SqliteAccountStore,
pub cookies: SqliteCookieStore,
pub codex: StdArc<CodexBackendClient>,
pub models: StdArc<ModelService>,
pub account_pool: StdArc<RuntimeAccountPoolService>,
pub token_refresher: StdArc<dyn TokenRefresher>,
pub oauth: oauth::AccountOAuthService,
@@ -58,6 +61,7 @@ impl AdminAccountService {
store: parts.store,
cookies: parts.cookies,
codex: parts.codex,
models: parts.models,
account_pool: parts.account_pool,
token_refresher: parts.token_refresher,
oauth: parts.oauth,
+12 -42
View File
@@ -7,11 +7,12 @@ use serde_json::{json, Value};
use tokio::sync::mpsc;
use crate::upstream::{
models::CodexModelInfo,
protocol::{
responses::CodexResponsesRequest,
sse::{encode_sse_event, parse_sse_events, SseEvent},
},
transport::{CodexModelCatalogClient, CodexModelCatalogRequest, CodexRequestContext},
transport::CodexRequestContext,
};
use super::{types::AdminAccountError, AdminAccountService};
@@ -29,7 +30,6 @@ impl AdminAccountService {
pub async fn account_models(
&self,
account_id: &str,
request_id: &str,
) -> Result<Vec<AccountModelOption>, AdminAccountError> {
let account = self
.store
@@ -37,21 +37,12 @@ impl AdminAccountService {
.await
.map_err(|_| AdminAccountError::Inspect)?
.ok_or(AdminAccountError::NotFound)?;
let token = account.access_token.expose_secret().to_string();
let request = CodexModelCatalogRequest {
access_token: &token,
account_id: account.account_id.as_deref(),
request_id,
installation_id: self.installation_id.as_deref(),
plan_type: account.plan_type.as_deref().unwrap_or("default"),
};
let models = self
.codex
.fetch_models(&request)
.await
.map_err(|error| AdminAccountError::FetchModels(error.to_string()))?
let plan_type = account.plan_type.as_deref().unwrap_or("default");
let catalog = self.models.catalog().await;
let models = catalog
.models_for_plan(plan_type)
.iter()
.filter_map(account_model_option)
.map(account_model_option)
.collect::<Vec<_>>();
if models.is_empty() {
return Err(AdminAccountError::NoModels);
@@ -145,32 +136,11 @@ impl AdminAccountService {
}
}
fn account_model_option(
entry: &crate::upstream::models::BackendModelEntry,
) -> Option<AccountModelOption> {
let id = first_non_empty([
entry.slug.as_deref(),
entry.id.as_deref(),
entry.name.as_deref(),
])?
.to_string();
let label = first_non_empty([
entry.display_name.as_deref(),
entry.title.as_deref(),
entry.name.as_deref(),
Some(id.as_str()),
])
.unwrap_or(&id)
.to_string();
Some(AccountModelOption { id, label })
}
fn first_non_empty<'a>(values: impl IntoIterator<Item = Option<&'a str>>) -> Option<&'a str> {
values
.into_iter()
.flatten()
.map(str::trim)
.find(|value| !value.is_empty())
fn account_model_option(model: &CodexModelInfo) -> AccountModelOption {
AccountModelOption {
id: model.id.clone(),
label: model.display_name.clone(),
}
}
fn test_responses_request(model: String) -> CodexResponsesRequest {
-2
View File
@@ -79,8 +79,6 @@ pub enum AdminAccountError {
StoreQuota,
#[error("failed to fetch quota: {0}")]
FetchQuota(String),
#[error("failed to fetch models: {0}")]
FetchModels(String),
#[error("no available models returned")]
NoModels,
#[error("invalid status: {0}")]
+12 -11
View File
@@ -190,10 +190,22 @@ impl Services {
stores.client_keys.clone(),
)));
let token_client = StdArc::new(default_openai_token_client(token_client_config(config)));
let upstream_client: StdArc<dyn CodexModelCatalogClient> = codex.clone();
let snapshot_store: StdArc<dyn ModelSnapshotStore> = StdArc::new(
crate::upstream::models::SqliteModelSnapshotStore::new(stores.accounts.pool().clone()),
);
let models = StdArc::new(ModelService::new(
crate::upstream::models::ModelConfig {
model_aliases: config.model_aliases.clone(),
},
Some(snapshot_store),
Some(upstream_client),
));
let admin_accounts = StdArc::new(AdminAccountService::new(AdminAccountServiceParts {
store: stores.accounts.clone(),
cookies: stores.cookies.clone(),
codex: codex.clone(),
models: models.clone(),
account_pool: account_pool.clone(),
token_refresher: token_client,
oauth: crate::admin::accounts::service::oauth::AccountOAuthService::new(
@@ -216,17 +228,6 @@ impl Services {
stores.session_affinity.clone(),
));
let upstream_client: StdArc<dyn CodexModelCatalogClient> = codex.clone();
let snapshot_store: StdArc<dyn ModelSnapshotStore> = StdArc::new(
crate::upstream::models::SqliteModelSnapshotStore::new(stores.accounts.pool().clone()),
);
let models = StdArc::new(ModelService::new(
crate::upstream::models::ModelConfig {
model_aliases: config.model_aliases.clone(),
},
Some(snapshot_store),
Some(upstream_client),
));
let cloudflare_recovery =
crate::proxy::dispatch::cloudflare::CloudflareRecovery::new(stores.cookies.clone());
let chat = StdArc::new(ChatDispatchService::new(
+13
View File
@@ -347,6 +347,19 @@ impl ModelCatalog {
self.models.clone()
}
/// 返回指定订阅计划可用的模型列表。
pub fn models_for_plan(&self, plan_type: &str) -> Vec<CodexModelInfo> {
self.models
.iter()
.filter(|model| {
self.model_plan_index
.get(&model.id)
.is_some_and(|plans| plans.iter().any(|plan| plan == plan_type))
})
.cloned()
.collect()
}
/// 按模型 ID 查询模型信息。
pub fn model_info(&self, model_id: &str) -> Option<CodexModelInfo> {
self.models
+48 -39
View File
@@ -35,24 +35,8 @@ fn test_events(body: &str) -> Vec<Value> {
}
#[tokio::test]
async fn account_models_should_return_upstream_models_only() {
async fn account_models_should_return_database_snapshot_models() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/codex/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"models": [
{
"slug": "gpt-5.5",
"display_name": "GPT 5.5"
},
{
"id": "gpt-5.4",
"title": "GPT 5.4"
}
]
})))
.mount(&server)
.await;
let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url(
"admin-account-models.sqlite",
91,
@@ -60,6 +44,7 @@ async fn account_models_should_return_upstream_models_only() {
)
.await;
seed_test_account(&pool).await;
seed_model_snapshot(&pool, "plus").await;
let response = app
.oneshot(
@@ -74,31 +59,55 @@ async fn account_models_should_return_upstream_models_only() {
.await
.unwrap();
let requests = server.received_requests().await.unwrap();
let upstream = requests
.iter()
.find(|request| request.url.path() == "/backend-api/codex/models")
.expect("models request should hit upstream");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
upstream
.headers
.get("authorization")
.and_then(|value| value.to_str().ok()),
Some("Bearer access-test")
);
assert_eq!(
upstream
.headers
.get("chatgpt-account-id")
.and_then(|value| value.to_str().ok()),
Some("chatgpt-test")
);
assert!(requests.is_empty(), "model list should be loaded from DB");
let body = response_json(response).await;
assert_eq!(body["data"]["models"][0]["id"], "gpt-5.5");
assert_eq!(body["data"]["models"][0]["label"], "GPT 5.5");
assert_eq!(body["data"]["models"][1]["id"], "gpt-5.4");
assert_eq!(body["data"]["models"][1]["label"], "GPT 5.4");
assert_eq!(body["data"]["models"][0]["id"], "gpt-5.4");
assert_eq!(body["data"]["models"][0]["label"], "GPT 5.4");
assert_eq!(body["data"]["models"][1]["id"], "gpt-5.5");
assert_eq!(body["data"]["models"][1]["label"], "GPT 5.5");
}
async fn seed_model_snapshot(pool: &SqlitePool, plan_type: &str) {
let models_json = json!([
{
"id": "gpt-5.5",
"displayName": "GPT 5.5",
"description": "Test model",
"isDefault": false,
"supportedReasoningEfforts": [{"reasoningEffort": "medium", "description": "medium"}],
"defaultReasoningEffort": "medium",
"inputModalities": ["text"],
"outputModalities": ["text"],
"supportsPersonality": false,
"upgrade": null,
"source": "test"
},
{
"id": "gpt-5.4",
"displayName": "GPT 5.4",
"description": "Test model",
"isDefault": false,
"supportedReasoningEfforts": [{"reasoningEffort": "medium", "description": "medium"}],
"defaultReasoningEffort": "medium",
"inputModalities": ["text"],
"outputModalities": ["text"],
"supportsPersonality": false,
"upgrade": null,
"source": "test"
}
])
.to_string();
sqlx::query(
"insert or replace into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?)",
)
.bind(plan_type)
.bind(models_json)
.bind("2026-06-29T00:00:00Z")
.execute(pool)
.await
.unwrap();
}
#[tokio::test]
+2
View File
@@ -105,6 +105,8 @@ fn model_catalog_should_merge_backend_snapshots_and_build_plan_allowlist() {
catalog.model_plan_allowlist().get("gpt-6").unwrap(),
&vec!["plus".to_string()]
);
assert_eq!(catalog.models_for_plan("plus")[0].id, "gpt-6");
assert!(catalog.models_for_plan("team").is_empty());
}
#[tokio::test]
@@ -194,7 +194,7 @@ export function useAccountConnectionTest() {
}))
connectionTestSelectedModel.value = connectionTestModelOptions.value[0]?.value || ''
if (!connectionTestSelectedModel.value) {
connectionTestError.value = '上游没有返回可测试模型'
connectionTestError.value = '没有可测试模型'
}
} catch (error: any) {
connectionTestError.value = error.message || '加载测试模型失败'
@@ -208,7 +208,7 @@ export function useAccountConnectionTest() {
async function handleTestConnection(account = testingAccount.value) {
if (!account?.id) return
if (!connectionTestSelectedModel.value) {
connectionTestError.value = '请先选择上游返回的测试模型'
connectionTestError.value = '请先选择测试模型'
return
}
if (testingConnectionIds.value.has(account.id)) return