From d461cb63253197e4394934cc0007d0a590290751 Mon Sep 17 00:00:00 2001 From: zyycn Date: Tue, 30 Jun 2026 00:00:24 +0800 Subject: [PATCH] fix: use model snapshots for account test models Co-authored-by: Codex --- src/admin/accounts/routes.rs | 7 +- src/admin/accounts/service/mod.rs | 4 + src/admin/accounts/service/testing.rs | 54 +++--------- src/admin/accounts/service/types.rs | 2 - src/runtime/services.rs | 23 ++--- src/upstream/models/mod.rs | 13 +++ tests/admin/accounts/testing.rs | 87 ++++++++++--------- tests/upstream/models/catalog.rs | 2 + .../composables/useAccountConnectionTest.ts | 4 +- 9 files changed, 95 insertions(+), 101 deletions(-) diff --git a/src/admin/accounts/routes.rs b/src/admin/accounts/routes.rs index 98d9d2336..90ade8d8c 100644 --- a/src/admin/accounts/routes.rs +++ b/src/admin/accounts/routes.rs @@ -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, Query(query): Query, - Extension(request_id): Extension, headers: HeaderMap, ) -> Result { - 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( diff --git a/src/admin/accounts/service/mod.rs b/src/admin/accounts/service/mod.rs index 33ae02be4..2dbe2a5ec 100644 --- a/src/admin/accounts/service/mod.rs +++ b/src/admin/accounts/service/mod.rs @@ -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, + pub(crate) models: StdArc, pub(crate) account_pool: StdArc, pub(crate) token_refresher: StdArc, pub(crate) oauth: oauth::AccountOAuthService, @@ -45,6 +47,7 @@ pub struct AdminAccountServiceParts { pub store: SqliteAccountStore, pub cookies: SqliteCookieStore, pub codex: StdArc, + pub models: StdArc, pub account_pool: StdArc, pub token_refresher: StdArc, 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, diff --git a/src/admin/accounts/service/testing.rs b/src/admin/accounts/service/testing.rs index c48d2785d..96f1c036f 100644 --- a/src/admin/accounts/service/testing.rs +++ b/src/admin/accounts/service/testing.rs @@ -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, 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::>(); if models.is_empty() { return Err(AdminAccountError::NoModels); @@ -145,32 +136,11 @@ impl AdminAccountService { } } -fn account_model_option( - entry: &crate::upstream::models::BackendModelEntry, -) -> Option { - 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>) -> 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 { diff --git a/src/admin/accounts/service/types.rs b/src/admin/accounts/service/types.rs index 179405877..bab6e3cb7 100644 --- a/src/admin/accounts/service/types.rs +++ b/src/admin/accounts/service/types.rs @@ -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}")] diff --git a/src/runtime/services.rs b/src/runtime/services.rs index c21bcac35..dfda25884 100644 --- a/src/runtime/services.rs +++ b/src/runtime/services.rs @@ -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 = codex.clone(); + let snapshot_store: StdArc = 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 = codex.clone(); - let snapshot_store: StdArc = 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( diff --git a/src/upstream/models/mod.rs b/src/upstream/models/mod.rs index e69ef1f5f..3b5d4f03b 100644 --- a/src/upstream/models/mod.rs +++ b/src/upstream/models/mod.rs @@ -347,6 +347,19 @@ impl ModelCatalog { self.models.clone() } + /// 返回指定订阅计划可用的模型列表。 + pub fn models_for_plan(&self, plan_type: &str) -> Vec { + 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 { self.models diff --git a/tests/admin/accounts/testing.rs b/tests/admin/accounts/testing.rs index a827f9523..ee30979eb 100644 --- a/tests/admin/accounts/testing.rs +++ b/tests/admin/accounts/testing.rs @@ -35,24 +35,8 @@ fn test_events(body: &str) -> Vec { } #[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] diff --git a/tests/upstream/models/catalog.rs b/tests/upstream/models/catalog.rs index 126541292..dd44d35c2 100644 --- a/tests/upstream/models/catalog.rs +++ b/tests/upstream/models/catalog.rs @@ -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] diff --git a/web/src/views/accounts/composables/useAccountConnectionTest.ts b/web/src/views/accounts/composables/useAccountConnectionTest.ts index 1d7528962..9c00029bc 100644 --- a/web/src/views/accounts/composables/useAccountConnectionTest.ts +++ b/web/src/views/accounts/composables/useAccountConnectionTest.ts @@ -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