mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
fix: use model snapshots for account test models
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user