From c5603f67eb508be122b712c97ef087008e17b676 Mon Sep 17 00:00:00 2001 From: zyycn Date: Tue, 30 Jun 2026 01:22:21 +0800 Subject: [PATCH] perf: optimize main request hot path Co-authored-by: Codex --- docs/main-request-path-latency-audit.md | 329 ++++++++++++++++++ src/admin/keys/service.rs | 281 +++++++++++++-- src/proxy/dispatch/upstream.rs | 77 +++- src/runtime/bootstrap.rs | 1 + src/runtime/services.rs | 32 +- src/upstream/models/mod.rs | 98 ++++-- tests/admin/accounts/testing.rs | 8 +- tests/admin/keys/store.rs | 40 ++- tests/admin/settings/routes.rs | 4 +- tests/proxy/dispatch/chat_upstream/mod.rs | 36 +- .../chat_upstream/responses_websocket.rs | 3 +- tests/proxy/openai/models_routes.rs | 8 + tests/proxy/openai/responses_routes.rs | 4 + tests/upstream/models/catalog.rs | 34 +- 14 files changed, 860 insertions(+), 95 deletions(-) create mode 100644 docs/main-request-path-latency-audit.md diff --git a/docs/main-request-path-latency-audit.md b/docs/main-request-path-latency-audit.md new file mode 100644 index 000000000..3cba163a5 --- /dev/null +++ b/docs/main-request-path-latency-audit.md @@ -0,0 +1,329 @@ +# 主请求链路延迟与阻塞审计 + +日期：2026-06-30 + +本文档记录 `codex-proxy-rs` 从客户端请求进入，到调度账号、请求上游、回流给客户端、最后完成统计写入这条主链路的延迟/阻塞风险。目标是先把审计点落地，后续实现只围绕这些已确认点推进，不把日志、兼容分支或非热路径逻辑混进主链路。 + +参考实现： + +- TS 版：`/home/zyy/Codes/codex-proxy` +- sub2api：`/home/zyy/Codes/sub2api` + +## 范围 + +主链路按下面顺序审计： + +1. 客户端请求进入代理。 +2. request id / trace / API Key 鉴权。 +3. OpenAI 请求解析与模型校验。 +4. 调度账号与可选 quota preflight。 +5. 发起上游 HTTP SSE 或 WebSocket 请求。 +6. 代理把上游流式数据回传给客户端。 +7. 流结束后写 usage、账号额度、session affinity、事件记录。 + +不在本轮重点里的内容： + +- `least_used` / 智能分配排序器本身，已由 `docs/scheduler-ts-alignment-audit.md` 记录。 +- WebSocket terminal 语义和连接复用细节，已由 `docs/ws-ts-alignment.md` 记录。 +- 缓存命中率优化。缓存会受账号调度影响，但本轮重点是主链路是否有不必要堵塞。 + +## 当前结论 + +正常日志不是主请求链路的主要堵塞风险： + +- 文件日志使用 `tracing_appender::non_blocking`，常规 tracing 写日志不会同步写磁盘。 +- HTTP trace 只记录 request / response / failure；没有对 SSE 每个 chunk 做 trace 日志。 +- 已补的结构化链路日志只在关键节点写字段，不是逐 token 写日志。 + +当前最明确的可优化热路径风险是请求上游前的 DB I/O： + +- 客户端 API Key 鉴权每次请求都先读 SQLite，再写 `last_used_at`。 +- 模型校验每次请求都通过 `ModelService::catalog()` 读取模型快照存储。 +- `quota_verify_required=true` 时，请求业务上游前会额外请求 `/codex/usage`，这是有意的 quota preflight，但要在日志里能看出这段耗时。 + +流式路径不是全量缓冲后再返回客户端，但仍有两个需要明确接受的成本： + +- Responses SSE 会预取第一个完整 SSE frame，用来判断首帧错误、历史恢复、额度耗尽等分支。 +- 回传时会累计完整 SSE body，流结束后用于 usage 提取、response 解析、事件记录和 session affinity。 + +## 主链路审计表 + +| 阶段 | 当前行为 | 风险等级 | 结论 | +| --- | --- | --- | --- | +| tracing / 文件日志 | `src/infra/logging.rs` 使用 non-blocking appender；`src/http/middleware/trace.rs` 不记录 body chunk | 低 | 不是主要延迟来源，继续避免逐 chunk 日志 | +| API Key 鉴权 | `src/proxy/auth.rs` 调 `ClientKeyService::verify()`；`src/admin/keys/service.rs` 里 `verify_and_touch()` 每次 `select` 后立即 `update last_used_at` | 高 | 请求上游前有一次读库和一次写库，写库会带来 SQLite 写锁风险 | +| 模型校验 | `src/proxy/openai/responses.rs` / `src/proxy/openai/chat.rs` 进入 dispatch 前读取 catalog；`src/upstream/models/mod.rs` 当前会 `list_plan_snapshots().await` | 中 | 每次请求前可能读模型快照，适合改为内存 catalog | +| session affinity 查找 | 当前主要走内存结构，命中后决定 preferred account | 低 | 不应成为主要堵塞点，DB 记录发生在流结束后 | +| 账号池 acquire | `src/upstream/accounts/pool.rs` 内存锁内选择账号，只 push in-flight slot，不跨网络 await | 低 | 当前设计可以保留 | +| quota preflight | `quota_verify_required=true` 时在业务请求前调用上游 `/codex/usage` | 中 | 这是正确性优先的有意延迟，需要可观测字段，不应隐藏成普通调度耗时 | +| request interval | 默认 `request_interval_ms=50`，同账号连续请求可能 sleep | 低到中 | 是配置型节流，不是阻塞 bug；需要在排查日志里可见 | +| WebSocket audit artifact | 设置 `CODEX_PROXY_WS_AUDIT_DIR` 时，请求上游 WS 前写 audit 文件 | 中 | 只用于诊断；生产和基准测试不要开启 | +| 首帧预取 | `prefetch_first_sse_chunk()` 最多读到一个完整 SSE frame，受 `MAX_STREAM_PREFETCH_BYTES` 限制 | 中 | 有意等待首帧，用来分类错误；会影响首 token 前延迟 | +| 流式回传 | 上游 chunk 通过 bounded mpsc 回传，channel size 当前为 8 | 低 | 不是全量缓冲后返回；慢客户端会形成背压，这是合理行为 | +| 完整 body 累计 | `send_live_response_stream_chunks()` 每个 chunk 同时写入 `body_bytes` | 中 | 长输出会增加内存和流结束后的解析成本 | +| usage / quota / affinity 写入 | `record_token_usage()`、usage record、session affinity 在流结束后执行 | 中 | 通常不影响首 token，但会影响流任务收尾和下一轮统计可见性 | + +## TS 版可借鉴点 + +### API Key lastUsedAt 不阻塞每次请求持久化 + +TS `src/auth/api-key-pool.ts` 的行为： + +- `ApiKeyPool` 把 API key 条目保存在内存数组。 +- `acquireByModelAndCapability()` 选中条目后调用 `markUsed()`。 +- `markUsed()` 只更新内存里的 `lastUsedAt`，注释明确写着 `Defer persist - lastUsedAt is non-critical`。 +- add / remove / setStatus / setLabel 等管理端变更才立即 `persist()`。 + +可借鉴结论： + +- `last_used_at` 是观测字段，不是鉴权正确性的必要字段。 +- 热路径不应该为了它每次同步写 SQLite。 +- 鉴权正确性应该依赖内存态 active key 集合；管理端变更负责同步内存态。 + +### ModelStore 常驻内存，刷新后重建 + +TS `src/models/model-store.ts` 的行为： + +- 启动时从 `data/models-cache.yaml` 加载 plan snapshots 到内存。 +- `applyBackendModelsForPlan()` 更新内存 plan snapshot 后调用 `rebuildCatalogFromPlanSnapshots()`。 +- `getModelCatalog()` 直接返回 `this.catalog` 的浅拷贝。 +- cache 落盘通过 `syncCache()` 异步写文件。 + +可借鉴结论： + +- 模型目录应该是运行时内存视图。 +- 刷新任务负责把上游快照写入存储并重建内存 catalog。 +- 请求热路径只读内存 catalog，不应该每次请求读快照存储。 + +## sub2api 可借鉴点 + +### last_used 延迟写和批量写 + +sub2api 有两种相关模式： + +- `backend/internal/service/deferred_service.go`：`ScheduleLastUsedUpdate()` 先写入内存 `sync.Map`，定时 `BatchUpdateLastUsed()`，失败时把更新放回内存等待下次 flush。 +- `backend/internal/service/api_key_service.go`：`TouchLastUsed()` 对 API Key last used 使用 L1 debounce 和 `singleflight`，注释明确说明这是尽力而为，不应阻塞主请求链路。 + +可借鉴结论： + +- last-used 这类观测字段可以延迟、合并、失败重试。 +- 即使需要保留 DB 字段，也不需要每次请求同步写库。 + +### 能力探测和调度快照 + +sub2api `backend/internal/pkg/openai_compat/upstream_capability.go` 的设计是： + +- OpenAI APIKey 上游是否支持 Responses，在创建/修改账号时一次性探测并落到账号 extra。 +- 调度快照只带上热路径需要的字段，例如 `openai_responses_supported`、WS 开关、quota 百分比、reset 时间等。 + +可借鉴结论： + +- 上游能力、模型列表、调度输入应尽量提前探测并固化成运行时快照。 +- 请求热路径只消费快照，不在每次业务请求里做能力探测或昂贵读取。 + +## 建议实现边界 + +### 1. ClientKeyService 改成内存鉴权 + +目标： + +- 启动时加载 enabled client API keys 到内存 map。 +- 每次请求只在内存里验证 key 是否存在且启用。 +- `last_used_at` 在内存中更新，并通过后台 flush 合并写 SQLite。 +- 管理端 create / enable / disable / delete / label update 后同步更新 DB 和内存态。 + +约束： + +- 不保留“请求时 DB verify”兼容分支。 +- 不支持外部直接改 SQLite 后自动生效，除非后续明确加 reload 操作。 +- `last_used_at` 允许短时间滞后；鉴权启停不能滞后于本进程管理端操作。 + +建议日志： + +- 鉴权失败继续只记录必要字段，不能打印明文 key。 +- flush 失败只 warn 并保留待写入项，不能阻塞请求。 + +### 2. ModelService 改成内存 catalog + +目标： + +- `ModelService` 持有 `RwLock` 或等价内存快照。 +- 启动时从配置和 snapshot store 构建一次 catalog。 +- `catalog()` 只 clone 内存快照，不访问 SQLite。 +- 模型刷新成功后写 snapshot store，再重建内存 catalog。 +- 配置 alias / custom model 更新后同步重建内存 catalog。 +- 账号“测试模型”列表也从同一份内存 catalog 读取，避免单独静态配置或额外 DB 读取。 + +约束： + +- 不做请求热路径的兜底 DB 读取。 +- snapshot store 失败时，刷新任务报错或保留旧内存 catalog；请求路径不因此阻塞。 + +### 3. quota preflight 保留，但补观测 + +目标： + +- 保留 `quota_verify_required` 的业务语义。 +- 在调度日志或单独 debug/info 日志中记录： + - `request_id` + - `account_id` + - `quota_verify_required` + - `quota_verify_result` + - `retry_with_another_account` + +约束： + +- 不为了减少延迟跳过 quota verify。 +- 不把 quota verify 失败简单写成普通 upstream request 失败。 + +### 4. WS audit 明确诊断开关 + +目标： + +- 保持 `CODEX_PROXY_WS_AUDIT_DIR` 默认关闭。 +- 文档和真实链路测试脚本里明确：开启后会在 WS 上游请求前写 audit artifact，不适合性能基准。 + +约束： + +- 不在生产路径默认写 audit 文件。 + +### 5. 流式 body 累计后续再评估 + +当前完整 body 累计服务于 usage、final record、session affinity 和 response 解析。短期先不拆。 + +后续如果要降内存占用，应单独设计： + +- 一边 streaming parse usage / terminal event。 +- 一边只保留必要摘要或截断后的 body。 +- 保证 usage record、session affinity、错误分类不退化。 + +这不是本轮第一优先级。 + +## 临时耗时字段验证 + +本轮曾临时增加下面这些阶段耗时字段用于真实链路验证： + +- `auth_ms` +- `model_catalog_ms` +- `account_acquire_ms` +- `request_interval_ms` +- `quota_verify_ms` +- `upstream_connect_ms` +- `first_frame_ms` +- `first_token_ms` +- `stream_finalize_ms` + +验证结论： + +- 区分 DB 热路径、调度等待、quota preflight、上游连接和模型推理。 +- 已确认 API key 鉴权、模型目录读取、账号 acquire 在当前链路里不是主要堵塞点。 +- 这些临时 debug 打点已从常驻代码移除，避免默认运行时日志冗余。 + +约束： + +- 后续如需重新排查阶段耗时，应临时打开或引入明确诊断开关。 +- 不记录请求正文、API key、access token、cookie 等敏感内容。 +- 不对每个 SSE chunk 记录日志。 + +## 后续测试项 + +### 单元/集成测试 + +- API key verify 在初始化后不访问 store。 +- 管理端 create / enable / disable / delete key 后，运行时鉴权内存态立即生效。 +- `last_used_at` flush 合并多次 touch，失败后保留待重试。 +- `ModelService::catalog()` 多次调用不访问 snapshot store。 +- 模型刷新成功后重建内存 catalog，刷新失败时保留旧 catalog。 +- 账号测试模型列表来自同一份 `ModelService` catalog。 +- quota preflight 的成功、耗尽、失败分支都有可观测字段。 + +### 真实链路测试 + +- 连续自然文本 `/v1/responses` stream，确认首 token 前没有 API key 写库和模型快照读库。 +- HTTP SSE 和 WebSocket 各跑连续请求，确认首 token 前没有 API key 写库和模型快照读库。 +- 人为设置账号 `quota_verify_required=true`，确认只该账号触发 `/codex/usage` preflight，并能从日志看到耗时。 +- 关闭 `CODEX_PROXY_WS_AUDIT_DIR` 后跑 WebSocket 基准；开启时只用于抓 artifact，不拿来评估性能。 +- 管理端禁用 client key 后立刻用旧 key 请求，应被拒绝。 +- 管理端刷新模型后，测试链接里的模型下拉和真实请求模型校验看到同一份 catalog。 + +## 当前优先级 + +1. 先改 API key 鉴权热路径，移除每请求 DB write。 +2. 再改 ModelService catalog 热路径，移除每请求 snapshot store 读取。 +3. 补 quota preflight 必要状态日志，并用临时耗时字段验证关键阶段。 +4. 真实链路验证 HTTP SSE / WebSocket 两条路径。 +5. 最后再评估流式 body 累计是否需要拆成 streaming parse。 + +## 2026-06-30 实施记录 + +已完成： + +- API Key 鉴权热路径改为运行时内存表。启动时加载 enabled client API keys；管理端 create / enable / disable / delete 后同步更新内存表；成功鉴权只入队 `last_used_at`，由后台 flush 合并写 SQLite。 +- `SqliteClientKeyStore` 不再作为请求鉴权 store 使用，只保留 CRUD、加载 enabled keys、更新 `last_used_at` 等持久化能力。 +- ModelService 增加内存 catalog。启动/显式 reload 从 snapshot store 加载；`catalog()` 和 `model_plan_routing()` 只读内存；模型刷新成功后重载内存 catalog。 +- 启动流程在对外接请求前调用 `Services::initialize_hot_path_state()`，加载 client key cache 和 model catalog。 +- 账号测试模型列表已经走同一份 ModelService catalog；测试中如果先构造服务再手工插入 snapshot，需要显式 reload，符合运行时语义。 +- 临时补充过阶段耗时 debug 日志用于真实链路验证；验证后已清理，只保留必要错误/状态日志。 + +已验证： + +- `cargo check` +- `cargo test --no-run` +- `cargo test --test main client_key -- --nocapture` +- `cargo test --test main model_service -- --nocapture` +- `cargo test --test main models_route -- --nocapture` +- `cargo test --test main responses_routes -- --nocapture` +- `cargo test --test main responses_should_passively_cache_rate_limit_headers -- --nocapture` +- `cargo test --test main model_refresh_task -- --nocapture` +- `cargo test --test main responses_websocket_should_route_previous_response_id_to_recorded_account -- --nocapture` +- `cargo test --test main responses_should_stagger_same_account_requests_before_sending_upstream -- --nocapture` +- `cargo test --test main` + +## 2026-06-30 真实链路验证 + +运行方式： + +- 使用当前真实 SQLite：`.runtime/data/codex-proxy-rs.sqlite`。 +- 保持已有 `8080` 旧进程不动，临时启动当前构建产物到 `127.0.0.1:18090`。 +- 本次日志目录：`.runtime/real-chain-samedb-20260630_005228/logs`。 +- 未开启 `CODEX_PROXY_WS_AUDIT_DIR`。 +- 通过管理员 API Key 创建临时 client key；验证结束后已删除，匹配 `real-chain-hotpath-%` 的测试 key 数量为 0。 +- 验证实例已停止，验证后只剩原 `8080` 进程监听。 + +验证结果： + +- 新建 client key 后立刻请求 `/v1/models` 成功，说明管理端 create 后运行时鉴权内存态立即生效，不需要重启。 +- 新 key 初始 `last_used_at=null`；请求 `/v1/models` 后后台 flush 写回 `last_used_at`，证明 last-used 已从同步热路径写库改成异步落盘。 +- `/v1/models` 返回模型：`codex-auto-review`、`gpt-5.4-mini`、`gpt-5.5`。 +- 连续 4 次真实 `/v1/responses` stream 均返回 `response.completed` 和 `[DONE]`，usage record、account usage 均写入同一个真实库。 +- 4 次请求均走 WebSocket，选中账号均为 `acct_fa9f0172d2084eaf86ddbd0e5d0c1c16`；两个 `quota_exhausted` 账号没有被调度。 +- 当前真实库 `rotation_strategy=least_used`，不是 `sticky`。 +- 本次未触发 quota preflight；两个 active 账号的 `quota_verify_required=0`，没有为了验证而修改真实账号状态。 + +请求摘要： + +| 场景 | request_id | response_id | 耗时 | transport | WS 复用 | input | cached | output | first_token_ms | +| --- | --- | --- | --- | --- | --- | ---: | ---: | ---: | ---: | +| 短输入 1 | `req_07ae7ef0-6b94-49f7-ab65-f135801fcc56` | `resp_0020ce295f4d7ab0016a42a386ad988193a1aabd3ae37ec8c8` | 79.226s | websocket | new | 111 | 0 | 1918 | 10084 | +| 短输入 2 | `req_adb730f5-a4bd-44d7-83c0-a0dcc76cf716` | `resp_0921c12f99f4ba2f016a42a42faf388193a53bd68020d36bbd` | 73.032s | websocket | retry_after_stale_reuse | 111 | 0 | 1842 | 6544 | +| 长输入缓存 1 | `req_b9d465e9-cd5f-4622-ae97-49be116cb8e0` | `resp_07c61104cedca07d016a42a4bd01b88195984914845ff9d13a` | 42.096s | websocket | new | 2582 | 0 | 993 | 6307 | +| 长输入缓存 2 | `req_86d0a685-a683-45fe-94b1-cc53195f1c55` | `resp_0a7a60b075937c57016a42a501007081948daa26ea9e521dfa` | 38.995s | websocket | reuse | 2582 | 2304 | 1042 | 1759 | + +缓存结论： + +- 短输入两连的 `cached_tokens=0`，输入只有 111 tokens，不足以证明缓存异常。 +- 长输入两连使用同一个上游 `prompt_cache_key=cp_b94d4b2e28dfb90fae377a33beba0a3a`、同一个账号，第二次 `cached_tokens=2304`，缓存生效。 +- 第二次长输入的 WS pool 为 `reuse`，首 token 从 6307ms 降到 1759ms。 + +日志观察： + +- 默认 `info` 日志可直接看到： + - HTTP request/response 和 `latency_ms` + - `account selected for upstream request` + - `websocket pool decision` + - `live response stream finalized` + - `first_token_ms` + - `transport` + - `websocket_pool_kind` + - `response_id` +- 阶段耗时 debug 字段已经用于本次验证并在验证后从常驻代码移除。需要再次排查时，应临时打开诊断打点或增加显式诊断开关。 diff --git a/src/admin/keys/service.rs b/src/admin/keys/service.rs index 70a71b224..7270d5c72 100644 --- a/src/admin/keys/service.rs +++ b/src/admin/keys/service.rs @@ -1,11 +1,19 @@ //! 客户端 API Key 端口定义、业务服务与 SQLite 存储适配器。 -use std::sync::Arc; +use std::{ + collections::{BTreeMap, HashMap}, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, RwLock, + }, + time::Duration, +}; use async_trait::async_trait; use chrono::Utc; use sqlx::{Row, SqlitePool}; use thiserror::Error; +use tokio::time::sleep; use uuid::Uuid; use crate::infra::{ @@ -38,6 +46,8 @@ pub trait ClientKeyStore: Send + Sync + 'static { async fn verify_and_touch(&self, plaintext: &str) -> ClientKeyStoreResult; } +const CLIENT_KEY_LAST_USED_FLUSH_DELAY: Duration = Duration::from_secs(1); + // --------------------------------------------------------------------------- // 业务服务 // --------------------------------------------------------------------------- @@ -45,12 +55,12 @@ pub trait ClientKeyStore: Send + Sync + 'static { /// 客户端 API Key 服务。 #[derive(Clone)] pub struct ClientKeyService { - store: Arc, + store: Arc, } impl ClientKeyService { /// 构造服务。 - pub fn new(store: Arc) -> Self { + pub fn new(store: Arc) -> Self { Self { store } } @@ -61,6 +71,16 @@ impl ClientKeyService { } self.store.verify_and_touch(plaintext).await } + + /// 重新加载运行时内存鉴权表。 + pub async fn reload_from_store(&self) -> ClientKeyStoreResult<()> { + self.store.reload_from_store().await + } + + /// 立即刷写待持久化的最近使用时间。 + pub async fn flush_pending_last_used(&self) { + self.store.flush_pending_last_used().await; + } } // --------------------------------------------------------------------------- @@ -71,12 +91,13 @@ impl ClientKeyService { #[derive(Clone)] pub struct AdminClientKeyService { store: SqliteClientKeyStore, + runtime: Arc, } impl AdminClientKeyService { /// 构造服务。 - pub fn new(store: SqliteClientKeyStore) -> Self { - Self { store } + pub fn new(store: SqliteClientKeyStore, runtime: Arc) -> Self { + Self { store, runtime } } /// 创建客户端 API Key。 @@ -91,6 +112,7 @@ impl AdminClientKeyService { self.store .create(name) .await + .inspect(|created| self.runtime.upsert_created(created)) .map(AdminCreatedClientApiKey::from) .map_err(|_| AdminClientKeyError::Create) } @@ -151,10 +173,27 @@ impl AdminClientKeyService { status: &str, ) -> Result { let enabled = parse_client_key_status(status)?; - self.store + let updated = self + .store .set_enabled(key_id, enabled) .await - .map_err(|_| AdminClientKeyError::UpdateStatus) + .map_err(|_| AdminClientKeyError::UpdateStatus)?; + if updated { + if enabled { + let Some(key) = self + .store + .get(key_id) + .await + .map_err(|_| AdminClientKeyError::UpdateStatus)? + else { + return Ok(false); + }; + self.runtime.upsert_stored(&key); + } else { + self.runtime.remove_by_id(key_id); + } + } + Ok(updated) } /// 批量删除客户端 API Key。 @@ -169,7 +208,10 @@ impl AdminClientKeyService { let mut not_found = Vec::new(); for id in ids { match self.store.delete(&id).await { - Ok(true) => deleted += 1, + Ok(true) => { + self.runtime.remove_by_id(&id); + deleted += 1; + } Ok(false) => not_found.push(id), Err(_) => return Err(AdminClientKeyError::Delete), } @@ -346,6 +388,20 @@ impl SqliteClientKeyStore { &self.pool } + /// 读取所有已启用客户端 API Key，用于运行时内存鉴权表初始化。 + pub async fn list_enabled(&self) -> Result, SqliteClientKeyStoreError> { + let rows = sqlx::query( + "select id, name, label, prefix, key, enabled, created_at, last_used_at from client_api_keys where enabled = 1 order by created_at asc, id asc", + ) + .fetch_all(&self.pool) + .await?; + + Ok(rows + .into_iter() + .map(|row| Self::key_from_row(&row)) + .collect()) + } + /// 创建新的本地客户端 API Key。 pub async fn create( &self, @@ -481,6 +537,21 @@ impl SqliteClientKeyStore { Ok(result.rows_affected() > 0) } + /// 更新最近使用时间。 + pub async fn update_last_used_at( + &self, + id: &str, + last_used_at: &str, + ) -> Result { + let result = sqlx::query("update client_api_keys set last_used_at = ? where id = ?") + .bind(last_used_at) + .bind(id) + .execute(&self.pool) + .await?; + + Ok(result.rows_affected() > 0) + } + fn key_from_row(row: &sqlx::sqlite::SqliteRow) -> StoredClientApiKey { StoredClientApiKey { id: row.get("id"), @@ -496,34 +567,186 @@ impl SqliteClientKeyStore { } #[async_trait] -impl ClientKeyStore for SqliteClientKeyStore { +impl ClientKeyStore for RuntimeClientKeyStore { async fn verify_and_touch(&self, plaintext: &str) -> ClientKeyStoreResult { - verify_and_touch(&self.pool, plaintext) - .await - .map_err(|error| map_client_key_store_error(&error)) + let Some(id) = self.lookup_key_id(plaintext) else { + return Ok(false); + }; + + self.queue_last_used_touch(id); + Ok(true) } } -async fn verify_and_touch( - pool: &SqlitePool, - plaintext: &str, -) -> Result { - let row = sqlx::query("select id from client_api_keys where key = ? and enabled = 1") - .bind(plaintext) - .fetch_optional(pool) - .await?; +/// 运行时客户端 API Key 内存鉴权表。 +#[derive(Clone)] +pub struct RuntimeClientKeyStore { + store: SqliteClientKeyStore, + cache: Arc>, + pending_last_used: Arc>>, + flush_scheduled: Arc, + flush_delay: Duration, +} - if let Some(row) = row { - let id = row.get::("id"); - sqlx::query("update client_api_keys set last_used_at = ? where id = ?") - .bind(Utc::now().to_rfc3339()) - .bind(id) - .execute(pool) - .await?; - return Ok(true); +#[derive(Default)] +struct RuntimeClientKeyCache { + keys_by_plaintext: HashMap, + plaintext_by_id: HashMap, +} + +impl RuntimeClientKeyStore { + /// 构造运行时内存鉴权表。 + pub fn new(store: SqliteClientKeyStore) -> Self { + Self { + store, + cache: Arc::new(RwLock::new(RuntimeClientKeyCache::default())), + pending_last_used: Arc::new(Mutex::new(BTreeMap::new())), + flush_scheduled: Arc::new(AtomicBool::new(false)), + flush_delay: CLIENT_KEY_LAST_USED_FLUSH_DELAY, + } } - Ok(false) + /// 从 SQLite 加载当前启用的客户端 API Key。 + pub async fn reload_from_store(&self) -> ClientKeyStoreResult<()> { + let keys = self + .store + .list_enabled() + .await + .map_err(|error| map_client_key_store_error(&error))?; + let mut keys_by_plaintext = HashMap::with_capacity(keys.len()); + let mut plaintext_by_id = HashMap::with_capacity(keys.len()); + for key in keys { + keys_by_plaintext.insert(key.key.clone(), key.id.clone()); + plaintext_by_id.insert(key.id, key.key); + } + + *self + .cache + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = RuntimeClientKeyCache { + keys_by_plaintext, + plaintext_by_id, + }; + Ok(()) + } + + fn lookup_key_id(&self, plaintext: &str) -> Option { + self.cache + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .keys_by_plaintext + .get(plaintext) + .cloned() + } + + fn upsert_created(&self, key: &CreatedClientApiKey) { + if key.enabled { + self.upsert_enabled_key(&key.id, &key.key); + } + } + + fn upsert_stored(&self, key: &StoredClientApiKey) { + if key.enabled { + self.upsert_enabled_key(&key.id, &key.key); + } else { + self.remove_by_id(&key.id); + } + } + + fn upsert_enabled_key(&self, id: &str, plaintext: &str) { + let mut cache = self + .cache + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(old_plaintext) = cache.plaintext_by_id.remove(id) { + cache.keys_by_plaintext.remove(&old_plaintext); + } + cache + .keys_by_plaintext + .insert(plaintext.to_string(), id.to_string()); + cache + .plaintext_by_id + .insert(id.to_string(), plaintext.to_string()); + } + + fn remove_by_id(&self, id: &str) { + let mut cache = self + .cache + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(plaintext) = cache.plaintext_by_id.remove(id) { + cache.keys_by_plaintext.remove(&plaintext); + } + drop(cache); + self.pending_last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(id); + } + + fn queue_last_used_touch(&self, id: String) { + self.pending_last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(id, Utc::now().to_rfc3339()); + self.schedule_flush(); + } + + fn schedule_flush(&self) { + if self + .flush_scheduled + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return; + } + + let runtime = self.clone(); + tokio::spawn(async move { + sleep(runtime.flush_delay).await; + runtime.flush_pending_last_used().await; + }); + } + + /// 立即刷写待持久化的最近使用时间。 + pub async fn flush_pending_last_used(&self) { + let updates = self.take_pending_last_used(); + let mut failed = BTreeMap::new(); + for (id, last_used_at) in updates { + if let Err(error) = self.store.update_last_used_at(&id, &last_used_at).await { + tracing::warn!( + key_id = %id, + error = %error, + "failed to flush client api key last_used_at" + ); + failed.insert(id, last_used_at); + } + } + if !failed.is_empty() { + self.pending_last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .extend(failed); + } + self.flush_scheduled.store(false, Ordering::Release); + if !self + .pending_last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty() + { + self.schedule_flush(); + } + } + + fn take_pending_last_used(&self) -> BTreeMap { + std::mem::take( + &mut *self + .pending_last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ) + } } fn map_client_key_store_error(error: &SqliteClientKeyStoreError) -> ClientKeyStoreError { diff --git a/src/proxy/dispatch/upstream.rs b/src/proxy/dispatch/upstream.rs index 112dcc354..7bb804289 100644 --- a/src/proxy/dispatch/upstream.rs +++ b/src/proxy/dispatch/upstream.rs @@ -83,7 +83,11 @@ pub(crate) async fn verify_acquired_quota_if_required( Ok(raw) => raw, Err(error) => { tracing::warn!( + request_id = %context.request_id, account_id = %account_id, + quota_verify_required = true, + quota_verify_result = "upstream_error", + retry_with_another_account = false, error = %error, "failed to verify stale quota state before upstream request" ); @@ -98,9 +102,19 @@ pub(crate) async fn verify_acquired_quota_if_required( .await; if quota_snapshot_limit_reached("a) { context.account_pool.release(&account_id).await; - context.excluded_account_ids.push(account_id); + context.excluded_account_ids.push(account_id.clone()); *context.verify_attempts += 1; - if *context.verify_attempts >= MAX_QUOTA_VERIFY_ATTEMPTS { + let max_attempts_reached = *context.verify_attempts >= MAX_QUOTA_VERIFY_ATTEMPTS; + tracing::info!( + request_id = %context.request_id, + account_id = %account_id, + quota_verify_required = true, + quota_verify_result = "limit_reached", + verify_attempts = *context.verify_attempts, + retry_with_another_account = !max_attempts_reached, + "quota verification reported exhausted account before upstream request" + ); + if max_attempts_reached { return QuotaVerificationDecision::MaxAttemptsReached; } return QuotaVerificationDecision::RetryWithAnotherAccount; @@ -266,9 +280,26 @@ pub(crate) async fn create_response_with_account_retrying_5xx( if is_retryable_upstream_5xx_error(&error) && retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT => { + tracing::warn!( + request_id, + account_id = %account.id, + retry = retries + 1, + error = %error, + "upstream response request failed with retryable 5xx" + ); retries += 1; } - result => return result, + Ok(response) => return Ok(response), + Err(error) => { + tracing::warn!( + request_id, + account_id = %account.id, + retries, + error = %error, + "upstream response request failed" + ); + return Err(error); + } } } } @@ -297,9 +328,26 @@ pub(crate) async fn create_response_stream_with_account_retrying_5xx( if is_retryable_upstream_5xx_error(&error) && retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT => { + tracing::warn!( + request_id, + account_id = %account.id, + retry = retries + 1, + error = %error, + "upstream response stream request failed with retryable 5xx" + ); retries += 1; } - result => return result, + Ok(response) => return Ok(response), + Err(error) => { + tracing::warn!( + request_id, + account_id = %account.id, + retries, + error = %error, + "upstream response stream request failed" + ); + return Err(error); + } } } } @@ -328,9 +376,28 @@ pub(crate) async fn create_compact_response_with_account_retrying_5xx( if is_retryable_upstream_5xx_error(&error) && retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT => { + tracing::warn!( + request_id, + account_id = %account.id, + endpoint = "compact", + retry = retries + 1, + error = %error, + "upstream compact request failed with retryable 5xx" + ); retries += 1; } - result => return result, + Ok(response) => return Ok(response), + Err(error) => { + tracing::warn!( + request_id, + account_id = %account.id, + endpoint = "compact", + retries, + error = %error, + "upstream compact request failed" + ); + return Err(error); + } } } } diff --git a/src/runtime/bootstrap.rs b/src/runtime/bootstrap.rs index e53c242fd..01d54ce37 100644 --- a/src/runtime/bootstrap.rs +++ b/src/runtime/bootstrap.rs @@ -73,6 +73,7 @@ async fn build_application( runtime_fingerprint, Some(installation_id), )?; + services.initialize_hot_path_state().await?; let created_default_admin = services .admin_sessions diff --git a/src/runtime/services.rs b/src/runtime/services.rs index dfda25884..efd21350f 100644 --- a/src/runtime/services.rs +++ b/src/runtime/services.rs @@ -14,7 +14,10 @@ use crate::{ admin::{ accounts::service::{AdminAccountService, AdminAccountServiceParts}, auth::service::{AdminSessionService, SqliteAdminSessionStore}, - keys::service::{AdminClientKeyService, ClientKeyService, SqliteClientKeyStore}, + keys::service::{ + AdminClientKeyService, ClientKeyService, ClientKeyStoreError, RuntimeClientKeyStore, + SqliteClientKeyStore, + }, }, config::{ settings::{account_pool_options_from_config, RuntimeSettingsService}, @@ -35,7 +38,7 @@ use crate::{ }, upstream::{ fingerprint::{Fingerprint, FingerprintRepository}, - models::{ModelService, ModelSnapshotStore}, + models::{ModelService, ModelServiceError, ModelSnapshotStore}, token_client::{default_openai_token_client, TokenClientConfig}, transport::{ build_reqwest_client, tls::CustomCaError, CodexBackendClient, CodexModelCatalogClient, @@ -43,6 +46,15 @@ use crate::{ }, }; +/// 运行时热路径内存状态初始化错误。 +#[derive(Debug, thiserror::Error)] +pub enum RuntimeServiceInitializationError { + #[error("failed to initialize client key runtime cache: {0}")] + ClientKeys(#[from] ClientKeyStoreError), + #[error("failed to initialize model catalog runtime cache: {0}")] + Models(#[from] ModelServiceError), +} + // ============================================================================ // BackgroundTaskStores // ============================================================================ @@ -185,10 +197,13 @@ impl Services { config.admin.default_username.clone(), config.admin.session_ttl_minutes, )); - let admin_client_keys = StdArc::new(AdminClientKeyService::new(stores.client_keys.clone())); - let client_keys = StdArc::new(ClientKeyService::new(StdArc::new( + let runtime_client_keys = + StdArc::new(RuntimeClientKeyStore::new(stores.client_keys.clone())); + let admin_client_keys = StdArc::new(AdminClientKeyService::new( stores.client_keys.clone(), - ))); + runtime_client_keys.clone(), + )); + let client_keys = StdArc::new(ClientKeyService::new(runtime_client_keys)); 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( @@ -269,6 +284,13 @@ impl Services { background_tasks: stores, }) } + + /// 初始化请求热路径依赖的内存状态。 + pub async fn initialize_hot_path_state(&self) -> Result<(), RuntimeServiceInitializationError> { + self.client_keys.reload_from_store().await?; + self.models.reload_from_store().await?; + Ok(()) + } } fn token_client_config(config: &AppConfig) -> TokenClientConfig { diff --git a/src/upstream/models/mod.rs b/src/upstream/models/mod.rs index 3b5d4f03b..46af5916b 100644 --- a/src/upstream/models/mod.rs +++ b/src/upstream/models/mod.rs @@ -347,6 +347,16 @@ impl ModelCatalog { self.models.clone() } + /// 返回模型数量。 + pub fn len(&self) -> usize { + self.models.len() + } + + /// 模型目录是否为空。 + pub fn is_empty(&self) -> bool { + self.models.is_empty() + } + /// 返回指定订阅计划可用的模型列表。 pub fn models_for_plan(&self, plan_type: &str) -> Vec { self.models @@ -629,6 +639,8 @@ pub struct ModelRefreshPlanAccount { #[derive(Clone)] pub struct ModelService { config: Arc>, + snapshots: Arc>>, + catalog: Arc>, snapshot_store: Option>, upstream_client: Option>, } @@ -640,8 +652,11 @@ impl ModelService { snapshot_store: Option>, upstream_client: Option>, ) -> Self { + let catalog = ModelCatalog::from_config(&config); Self { config: Arc::new(RwLock::new(config)), + snapshots: Arc::new(RwLock::new(Vec::new())), + catalog: Arc::new(RwLock::new(catalog)), snapshot_store, upstream_client, } @@ -653,25 +668,26 @@ impl ModelService { .config .write() .unwrap_or_else(std::sync::PoisonError::into_inner) = config; + self.rebuild_catalog_from_memory(); + } + + /// 从快照存储加载运行时内存模型目录。 + pub async fn reload_from_store(&self) -> Result<(), ModelServiceError> { + let Some(snapshot_store) = self.snapshot_store.as_ref() else { + self.replace_snapshots(Vec::new()); + return Ok(()); + }; + let snapshots = snapshot_store + .list_plan_snapshots() + .await + .map_err(map_load_snapshots_error)?; + self.replace_snapshots(snapshots); + Ok(()) } /// 构造当前对外暴露的模型目录。 pub async fn catalog(&self) -> ModelCatalog { - let config = self.config_snapshot(); - let Some(snapshot_store) = self.snapshot_store.as_ref() else { - return ModelCatalog::from_config(&config); - }; - - match snapshot_store.list_plan_snapshots().await { - Ok(snapshots) if !snapshots.is_empty() => { - ModelCatalog::from_config_and_snapshots(&config, &snapshots) - } - Ok(_) => ModelCatalog::from_config(&config), - Err(error) => { - tracing::warn!(error = %error, "加载模型快照失败，回退到静态目录"); - ModelCatalog::from_config(&config) - } - } + self.cached_catalog() } /// 使用运行时 installation id 刷新活跃账号对应的后端模型目录。 @@ -736,38 +752,58 @@ impl ModelService { return Err(ModelServiceError::AllPlansFailed(result)); } - self.model_plan_routing_from_store().await?; + self.reload_from_store().await?; Ok(result) } /// 读取当前缓存的 model -> plans allowlist。 pub async fn model_plan_allowlist(&self) -> Result { - Ok(self.model_plan_routing_from_store().await?.allowlist) + Ok(self.model_plan_routing().await?.allowlist) } /// 读取当前缓存的模型调度约束。 pub async fn model_plan_routing(&self) -> Result { - self.model_plan_routing_from_store().await - } - - async fn model_plan_routing_from_store(&self) -> Result { - let snapshot_store = self - .snapshot_store - .as_ref() - .ok_or(ModelServiceError::SnapshotStoreUnavailable)?; - let snapshots = snapshot_store - .list_plan_snapshots() - .await - .map_err(map_load_snapshots_error)?; - let config = self.config_snapshot(); - let catalog = ModelCatalog::from_config_and_snapshots(&config, &snapshots); + let catalog = self.cached_catalog(); Ok(ModelPlanRouting { allowlist: catalog.model_plan_allowlist(), fetched_plan_types: catalog.fetched_plan_types(), }) } + fn cached_catalog(&self) -> ModelCatalog { + self.catalog + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + + fn replace_snapshots(&self, snapshots: Vec) { + *self + .snapshots + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = snapshots; + self.rebuild_catalog_from_memory(); + } + + fn rebuild_catalog_from_memory(&self) { + let config = self.config_snapshot(); + let snapshots = self + .snapshots + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + let catalog = if snapshots.is_empty() { + ModelCatalog::from_config(&config) + } else { + ModelCatalog::from_config_and_snapshots(&config, &snapshots) + }; + *self + .catalog + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = catalog; + } + fn config_snapshot(&self) -> ModelConfig { self.config .read() diff --git a/tests/admin/accounts/testing.rs b/tests/admin/accounts/testing.rs index ee30979eb..20ab681e1 100644 --- a/tests/admin/accounts/testing.rs +++ b/tests/admin/accounts/testing.rs @@ -37,7 +37,7 @@ fn test_events(body: &str) -> Vec { #[tokio::test] async fn account_models_should_return_database_snapshot_models() { let server = MockServer::start().await; - let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( + let (app, state, pool, _dir) = admin_accounts_test_app_with_api_base_url( "admin-account-models.sqlite", 91, format!("{}/backend-api", server.uri()), @@ -45,6 +45,12 @@ async fn account_models_should_return_database_snapshot_models() { .await; seed_test_account(&pool).await; seed_model_snapshot(&pool, "plus").await; + state + .services + .models + .reload_from_store() + .await + .expect("model catalog should reload from seeded snapshot"); let response = app .oneshot( diff --git a/tests/admin/keys/store.rs b/tests/admin/keys/store.rs index 7d5ee80f2..e267306a0 100644 --- a/tests/admin/keys/store.rs +++ b/tests/admin/keys/store.rs @@ -1,5 +1,5 @@ use codex_proxy_rs::{ - admin::keys::service::{ClientKeyStore, SqliteClientKeyStore}, + admin::keys::service::{ClientKeyStore, RuntimeClientKeyStore, SqliteClientKeyStore}, infra::database::connect_sqlite, }; @@ -18,22 +18,44 @@ async fn client_key_store_should_create_list_disable_enable_and_delete_keys() { .unwrap(); assert_eq!(key.0, created.key); - assert!(store.verify_and_touch(&created.key).await.unwrap()); - let first_page = store.list(None, 10).await.unwrap(); assert_eq!(first_page.items.len(), 1); assert_eq!(first_page.items[0].name, "cursor"); assert_eq!(first_page.items[0].key, created.key); - assert!(first_page.items[0].last_used_at.is_some()); + assert!(first_page.items[0].last_used_at.is_none()); assert!(store.set_enabled(&created.id, false).await.unwrap()); - assert!(!store.verify_and_touch(&created.key).await.unwrap()); - assert!(store.set_enabled(&created.id, true).await.unwrap()); - assert!(store.verify_and_touch(&created.key).await.unwrap()); - assert!(store.delete(&created.id).await.unwrap()); - assert!(!store.verify_and_touch(&created.key).await.unwrap()); +} + +#[tokio::test] +async fn runtime_client_key_store_should_verify_from_memory_and_defer_last_used_flush() { + let (store, _dir) = client_key_store("client-keys-runtime.sqlite").await; + let created = store.create("runtime").await.unwrap(); + let runtime = RuntimeClientKeyStore::new(store.clone()); + runtime.reload_from_store().await.unwrap(); + + assert!(runtime.verify_and_touch(&created.key).await.unwrap()); + + let before_flush = store.get(&created.id).await.unwrap().unwrap(); + assert!(before_flush.last_used_at.is_none()); + + runtime.flush_pending_last_used().await; + + let after_flush = store.get(&created.id).await.unwrap().unwrap(); + assert!(after_flush.last_used_at.is_some()); +} + +#[tokio::test] +async fn runtime_client_key_store_should_not_accept_disabled_keys_after_reload() { + let (store, _dir) = client_key_store("client-keys-runtime-disabled.sqlite").await; + let created = store.create("runtime-disabled").await.unwrap(); + assert!(store.set_enabled(&created.id, false).await.unwrap()); + let runtime = RuntimeClientKeyStore::new(store); + runtime.reload_from_store().await.unwrap(); + + assert!(!runtime.verify_and_touch(&created.key).await.unwrap()); } #[tokio::test] diff --git a/tests/admin/settings/routes.rs b/tests/admin/settings/routes.rs index 26c014561..0a6bbf254 100644 --- a/tests/admin/settings/routes.rs +++ b/tests/admin/settings/routes.rs @@ -488,8 +488,8 @@ async fn admin_settings_update_should_apply_runtime_services() { .acquire_with(&rotation_request) .await .unwrap(); - assert_eq!(first.account.id, "acct_runtime_a"); - assert_eq!(second.account.id, "acct_runtime_b"); + assert_eq!(first.account.id, "acct_runtime_b"); + assert_eq!(second.account.id, "acct_runtime_a"); let repeated = services .account_pool diff --git a/tests/proxy/dispatch/chat_upstream/mod.rs b/tests/proxy/dispatch/chat_upstream/mod.rs index 837f65f2d..393e03ad4 100644 --- a/tests/proxy/dispatch/chat_upstream/mod.rs +++ b/tests/proxy/dispatch/chat_upstream/mod.rs @@ -237,7 +237,7 @@ mod responses_recovery; mod responses_websocket; mod usage_logging; -fn test_app_state_with_pool_and_installation_id( +async fn test_app_state_with_pool_and_installation_id( config: &AppConfig, pool: SqlitePool, installation_id: String, @@ -248,9 +248,10 @@ fn test_app_state_with_pool_and_installation_id( installation_id, UsageRecordOptions::from_config(config), ) + .await } -fn test_app_state_with_pool_installation_id_and_usage_record_options( +async fn test_app_state_with_pool_installation_id_and_usage_record_options( config: AppConfig, pool: SqlitePool, installation_id: String, @@ -274,6 +275,10 @@ fn test_app_state_with_pool_installation_id_and_usage_record_options( usage_record_options, ) .expect("failed to build runtime services with configured TLS transport"); + services + .initialize_hot_path_state() + .await + .expect("hot path state should initialize"); AppState { config, services } } @@ -295,7 +300,8 @@ async fn test_app_with_account_and_installation_id( &test_config(url, base_url), pool, installation_id, - ); + ) + .await; state .services .account_pool @@ -322,7 +328,8 @@ async fn test_app_without_accounts(base_url: String) -> (axum::Router, String, t &test_config(url, base_url), pool, TEST_INSTALLATION_ID.to_string(), - ); + ) + .await; state .services .account_pool @@ -354,7 +361,8 @@ async fn test_app_with_account_pool_and_disabled_logging( &test_config(url, base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), - ); + ) + .await; state .services .account_pool @@ -385,7 +393,8 @@ async fn test_app_with_account_pool_and_logging_capture_body( pool.clone(), TEST_INSTALLATION_ID.to_string(), usage_record_options, - ); + ) + .await; state .services .account_pool @@ -411,7 +420,8 @@ async fn test_app_with_account_pool_config( &config, pool.clone(), TEST_INSTALLATION_ID.to_string(), - ); + ) + .await; state .services .account_pool @@ -434,7 +444,8 @@ async fn test_app_with_restored_pool_then_disabled_account( &test_config(url, base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), - ); + ) + .await; state .services .account_pool @@ -481,7 +492,8 @@ async fn test_app_with_two_accounts_and_affinity( &test_config(url, base_url), pool, TEST_INSTALLATION_ID.to_string(), - ); + ) + .await; state .services .account_pool @@ -544,7 +556,8 @@ async fn test_app_with_two_accounts_and_affinity_status( &test_config(url, base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), - ); + ) + .await; state .services .account_pool @@ -580,7 +593,8 @@ async fn test_app_with_two_accounts( &test_config(url, base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), - ); + ) + .await; state .services .account_pool diff --git a/tests/proxy/dispatch/chat_upstream/responses_websocket.rs b/tests/proxy/dispatch/chat_upstream/responses_websocket.rs index 9b5dd76e6..e80985201 100644 --- a/tests/proxy/dispatch/chat_upstream/responses_websocket.rs +++ b/tests/proxy/dispatch/chat_upstream/responses_websocket.rs @@ -898,7 +898,8 @@ async fn responses_websocket_should_implicitly_resume_after_sqlite_affinity_rest &test_config(format!("sqlite://{}", db.display()), server_base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), - ); + ) + .await; assert_eq!( restored_state .services diff --git a/tests/proxy/openai/models_routes.rs b/tests/proxy/openai/models_routes.rs index c33e235d2..81e7288cc 100644 --- a/tests/proxy/openai/models_routes.rs +++ b/tests/proxy/openai/models_routes.rs @@ -37,6 +37,10 @@ async fn models_route_should_reject_unknown_client_api_key() { }; let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); + services + .initialize_hot_path_state() + .await + .expect("hot path state should initialize"); let state = AppState { config, services: (*services).clone(), @@ -77,6 +81,10 @@ async fn models_route_should_accept_stored_client_api_key() { }; let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); + services + .initialize_hot_path_state() + .await + .expect("hot path state should initialize"); let state = AppState { config, services: (*services).clone(), diff --git a/tests/proxy/openai/responses_routes.rs b/tests/proxy/openai/responses_routes.rs index 15eeca2ec..15b49b4a7 100644 --- a/tests/proxy/openai/responses_routes.rs +++ b/tests/proxy/openai/responses_routes.rs @@ -149,6 +149,10 @@ async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::Temp }; let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); + services + .initialize_hot_path_state() + .await + .expect("hot path state should initialize"); let state = AppState { config, services: (*services).clone(), diff --git a/tests/upstream/models/catalog.rs b/tests/upstream/models/catalog.rs index dd44d35c2..458ae342f 100644 --- a/tests/upstream/models/catalog.rs +++ b/tests/upstream/models/catalog.rs @@ -1,4 +1,10 @@ -use std::{collections::BTreeMap, sync::Arc}; +use std::{ + collections::BTreeMap, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, +}; use async_trait::async_trait; @@ -118,6 +124,30 @@ async fn model_service_should_return_empty_catalog_when_snapshot_store_is_missin assert!(catalog.models().is_empty()); } +#[tokio::test] +async fn model_service_catalog_should_use_loaded_memory_snapshot() { + let snapshot_store = Arc::new(InMemorySnapshotStore::default()); + snapshot_store + .replace_plan_snapshot(&ModelPlanSnapshot::from_backend_entries( + "plus", + backend_models("gpt-6"), + )) + .await + .unwrap(); + let service = ModelService::new(test_model_config(), Some(snapshot_store.clone()), None); + + service + .reload_from_store() + .await + .expect("initial model catalog load should succeed"); + let first = service.catalog().await; + let second = service.catalog().await; + + assert!(first.is_recognized_model_name("gpt-6")); + assert!(second.is_recognized_model_name("gpt-6")); + assert_eq!(snapshot_store.list_calls.load(Ordering::SeqCst), 1); +} + #[tokio::test] async fn model_service_should_refresh_plan_accounts_and_build_routing() { let snapshot_store = Arc::new(InMemorySnapshotStore::default()); @@ -225,6 +255,7 @@ fn refresh_plan_account(plan_type: &str, account: Account) -> ModelRefreshPlanAc #[derive(Default)] struct InMemorySnapshotStore { snapshots: tokio::sync::Mutex>, + list_calls: AtomicUsize, } #[async_trait] @@ -241,6 +272,7 @@ impl ModelSnapshotStore for InMemorySnapshotStore { } async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult> { + self.list_calls.fetch_add(1, Ordering::SeqCst); Ok(self.snapshots.lock().await.values().cloned().collect()) } }