mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
feat(gateway): 重构 usage 数据层、迁移系统与系统导入
数据库迁移: - 引入 baseline v2 bootstrap,空库首次启动自动初始化 - 服务启动不再自动执行迁移,需显式 `--migrate` 运行 - 新增 pending migration 检测,schema 落后时拒绝启动 Usage 数据层: - usage body 存储外部化为独立 blob 表 - 新增 HTTP audit 表拆分存储请求/响应头与 body ref - 后台清理任务支持 legacy body ref 元数据迁移 - usage runtime 写入迁移到专用 tokio runtime(独立线程池, 8MB 栈) 系统导入/导出: - 支持用户、API Keys、钱包数据的完整导入 - 兼容 legacy 与 v1.3+ 两种导出格式 其他改进: - executor outcome 增加 runtime miss 诊断上下文 - 主 tokio runtime 栈大小调整为 8MB - 前端 provider 管理支持 base URL 配置 - dev.sh 支持 --migrate 参数
This commit is contained in:
@@ -652,6 +652,10 @@ async fn gateway_surfaces_local_execution_runtime_miss_reason_when_all_openai_ch
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"没有可用的提供商支持模型 gpt-5 的同步请求"
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-openai-chat-local-miss-123")
|
||||
|
||||
@@ -7,6 +7,7 @@ use super::{
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER,
|
||||
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
@@ -451,6 +452,131 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_surfaces_candidate_list_empty_reason_for_claude_chat_runtime_miss() {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:chat"])),
|
||||
Some(serde_json::json!(["claude-sonnet-4-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:chat"])),
|
||||
Some(serde_json::json!(["claude-sonnet-4-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/v1/messages",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
let execution_runtime = Router::new();
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-claude-chat-empty")),
|
||||
sample_auth_snapshot("api-key-claude-empty-1", "user-claude-empty-1"),
|
||||
)]));
|
||||
let candidate_selection_repository = Arc::new(
|
||||
InMemoryMinimalCandidateSelectionReadRepository::seed(vec![]),
|
||||
);
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages?beta=true"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("x-api-key", "sk-client-claude-chat-empty")
|
||||
.header("anthropic-version", "2023-06-01")
|
||||
.header(TRACE_ID_HEADER, "trace-claude-chat-empty-123")
|
||||
.body("{\"model\":\"claude-sonnet-4-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("candidate_list_empty")
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"没有可用的提供商支持模型 claude-sonnet-4-5 的同步请求"
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-claude-chat-empty-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert!(stored_candidates.is_empty());
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_claude_chat_error_for_local_sync_failure() {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
|
||||
@@ -7,6 +7,7 @@ use super::{
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER,
|
||||
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
@@ -732,3 +733,247 @@ async fn gateway_returns_claude_cli_error_for_local_sync_failure() {
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_marks_claude_cli_cross_format_runtime_miss_when_format_conversion_is_disabled() {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["gpt-5.4"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
None,
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["gpt-5.4"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-claude-cli-openai-local-miss-1".to_string(),
|
||||
provider_name: "RightCode".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-claude-cli-openai-local-miss-1".to_string(),
|
||||
endpoint_api_format: "openai:cli".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-claude-cli-openai-local-miss-1".to_string(),
|
||||
key_name: "codex".to_string(),
|
||||
key_auth_type: "bearer".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:cli".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:cli": 1})),
|
||||
model_id: "model-claude-cli-openai-local-miss-1".to_string(),
|
||||
global_model_id: "global-model-claude-cli-openai-local-miss-1".to_string(),
|
||||
global_model_name: "gpt-5.4".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-5.4".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5.4".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:cli".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-claude-cli-openai-local-miss-1".to_string(),
|
||||
"RightCode".to_string(),
|
||||
Some("https://right.codes".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-claude-cli-openai-local-miss-1".to_string(),
|
||||
"provider-claude-cli-openai-local-miss-1".to_string(),
|
||||
"openai:cli".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://right.codes/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some("/v1/messages".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-claude-cli-openai-local-miss-1".to_string(),
|
||||
"provider-claude-cli-openai-local-miss-1".to_string(),
|
||||
"codex".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["openai:cli"])),
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"sk-upstream-openai-cli-local-miss",
|
||||
)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"openai:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/v1/messages",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
let execution_runtime = Router::new();
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-claude-cli-openai-local-miss")),
|
||||
sample_auth_snapshot(
|
||||
"api-key-claude-cli-openai-local-miss-1",
|
||||
"user-claude-cli-openai-local-miss-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages?beta=true"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-claude-cli-openai-local-miss",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-claude-cli-openai-local-miss-123")
|
||||
.body("{\"model\":\"gpt-5.4\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("all_candidates_skipped")
|
||||
);
|
||||
let response_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(response_json["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
response_json["error"]["message"],
|
||||
"没有可用的提供商支持模型 gpt-5.4 的同步请求"
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-claude-cli-openai-local-miss-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Skipped);
|
||||
assert_eq!(
|
||||
stored_candidates[0].skip_reason.as_deref(),
|
||||
Some("format_conversion_disabled")
|
||||
);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
@@ -813,7 +813,7 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
"enable_format_conversion": false,
|
||||
"config": {"provider_ops": {"architecture_id": "cubence"}},
|
||||
"claude_code_advanced": {"pool_size": 2},
|
||||
"pool_advanced": {"enabled": true},
|
||||
"pool_advanced": {},
|
||||
"failover_rules": {"strategy": "ordered"},
|
||||
"proxy": {"url": "https://proxy.example"}
|
||||
}))
|
||||
@@ -839,7 +839,7 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
assert_eq!(payload["stream_first_byte_timeout"], 11.0);
|
||||
assert_eq!(payload["proxy"], json!({"url": "https://proxy.example"}));
|
||||
assert_eq!(payload["claude_code_advanced"], json!({"pool_size": 2}));
|
||||
assert_eq!(payload["pool_advanced"], json!({"enabled": true}));
|
||||
assert_eq!(payload["pool_advanced"], json!({}));
|
||||
assert_eq!(payload["failover_rules"], json!({"strategy": "ordered"}));
|
||||
assert_eq!(payload["ops_configured"], true);
|
||||
assert_eq!(payload["ops_architecture_id"], "cubence");
|
||||
@@ -895,7 +895,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
"website": "codex.example",
|
||||
"keep_priority_on_conversion": true,
|
||||
"max_retries": 7,
|
||||
"pool_advanced": {"enabled": true},
|
||||
"pool_advanced": {},
|
||||
"failover_rules": {"strategy": "ordered"},
|
||||
"proxy": {"url": "https://proxy.example"}
|
||||
}))
|
||||
@@ -929,6 +929,14 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() {
|
||||
assert!(created.enable_format_conversion);
|
||||
assert_eq!(created.max_retries, Some(7));
|
||||
assert_eq!(created.keep_priority_on_conversion, true);
|
||||
assert_eq!(
|
||||
created
|
||||
.config
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("pool_advanced"))
|
||||
.cloned(),
|
||||
Some(json!({}))
|
||||
);
|
||||
|
||||
let endpoints = provider_catalog_repository
|
||||
.list_endpoints_by_provider_ids(std::slice::from_ref(&created.id))
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository,
|
||||
InMemoryAuthApiKeySnapshotRepository, StandaloneApiKeyExportListQuery,
|
||||
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||
@@ -9,6 +11,7 @@ use aether_data::repository::users::{
|
||||
InMemoryUserReadRepository, StoredUserAuthRecord, StoredUserSummary,
|
||||
};
|
||||
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||
use async_trait::async_trait;
|
||||
use axum::body::Body;
|
||||
use axum::routing::{any, get};
|
||||
use axum::{extract::Request, Router};
|
||||
@@ -153,6 +156,103 @@ fn sample_api_key_snapshot(
|
||||
snapshot
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct PartialListAuthApiKeyRepository {
|
||||
lookup: InMemoryAuthApiKeySnapshotRepository,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthApiKeyReadRepository for PartialListAuthApiKeyRepository {
|
||||
async fn find_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, aether_data::DataLayerError> {
|
||||
self.lookup.find_api_key_snapshot(key).await
|
||||
}
|
||||
|
||||
async fn list_api_key_snapshots_by_ids(
|
||||
&self,
|
||||
_api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeySnapshot>, aether_data::DataLayerError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn list_export_api_keys_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, aether_data::DataLayerError> {
|
||||
self.lookup.list_export_api_keys_by_user_ids(user_ids).await
|
||||
}
|
||||
|
||||
async fn list_export_api_keys_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, aether_data::DataLayerError> {
|
||||
self.lookup.list_export_api_keys_by_ids(api_key_ids).await
|
||||
}
|
||||
|
||||
async fn list_export_standalone_api_keys_page(
|
||||
&self,
|
||||
query: &StandaloneApiKeyExportListQuery,
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, aether_data::DataLayerError> {
|
||||
self.lookup
|
||||
.list_export_standalone_api_keys_page(query)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn count_export_standalone_api_keys(
|
||||
&self,
|
||||
is_active: Option<bool>,
|
||||
) -> Result<u64, aether_data::DataLayerError> {
|
||||
self.lookup
|
||||
.count_export_standalone_api_keys(is_active)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn summarize_export_api_keys_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, aether_data::DataLayerError> {
|
||||
self.lookup
|
||||
.summarize_export_api_keys_by_user_ids(user_ids, now_unix_secs)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn summarize_export_non_standalone_api_keys(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, aether_data::DataLayerError> {
|
||||
self.lookup
|
||||
.summarize_export_non_standalone_api_keys(now_unix_secs)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn summarize_export_standalone_api_keys(
|
||||
&self,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<AuthApiKeyExportSummary, aether_data::DataLayerError> {
|
||||
self.lookup
|
||||
.summarize_export_standalone_api_keys(now_unix_secs)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn find_export_standalone_api_key_by_id(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, aether_data::DataLayerError> {
|
||||
self.lookup
|
||||
.find_export_standalone_api_key_by_id(api_key_id)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_export_standalone_api_keys(
|
||||
&self,
|
||||
) -> Result<Vec<StoredAuthApiKeyExportRecord>, aether_data::DataLayerError> {
|
||||
self.lookup.list_export_standalone_api_keys().await
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_stats_provider_quota_usage_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
@@ -1102,6 +1202,63 @@ async fn gateway_handles_admin_stats_leaderboard_api_keys_locally_without_auth_s
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_stats_leaderboard_api_keys_with_auth_snapshot_single_lookup_fallback(
|
||||
) {
|
||||
let (_upstream_url, upstream_hits, upstream_handle) =
|
||||
start_stats_upstream("/api/admin/stats/leaderboard/api-keys").await;
|
||||
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![sample_usage_row(
|
||||
"usage-key-partial-list",
|
||||
"req-key-partial-list",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("legacy-key"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
80,
|
||||
20,
|
||||
0.3,
|
||||
0.3,
|
||||
DAY_1_UNIX_SECS,
|
||||
)]));
|
||||
let auth_repository = Arc::new(PartialListAuthApiKeyRepository {
|
||||
lookup: InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
None,
|
||||
sample_api_key_snapshot("key-1", "user-1", "fresh-key"),
|
||||
)]),
|
||||
});
|
||||
|
||||
let data_state = GatewayDataState::with_usage_reader_for_tests(usage_repository)
|
||||
.with_auth_api_key_reader(auth_repository);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(
|
||||
reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/stats/leaderboard/api-keys?start_date=2024-03-21&end_date=2024-03-21&metric=cost&order=desc&tz_offset_minutes=0"
|
||||
)),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["items"][0]["id"], "key-1");
|
||||
assert_eq!(payload["items"][0]["name"], "fresh-key");
|
||||
assert_eq!(payload["items"][0]["value"], 0.3);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_stats_leaderboard_api_keys_locally_without_usage_reader() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
@@ -1206,21 +1363,23 @@ async fn gateway_handles_admin_stats_leaderboard_users_locally_without_user_read
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_stats_upstream("/api/admin/stats/leaderboard/users").await;
|
||||
|
||||
let mut usage_user = sample_usage_row(
|
||||
"usage-user-a-fallback",
|
||||
"req-user-a-fallback",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary-key"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
60,
|
||||
20,
|
||||
0.4,
|
||||
0.4,
|
||||
DAY_1_UNIX_SECS,
|
||||
);
|
||||
usage_user.username = Some("stale-alice".to_string());
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_usage_row(
|
||||
"usage-user-a-fallback",
|
||||
"req-user-a-fallback",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary-key"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
60,
|
||||
20,
|
||||
0.4,
|
||||
0.4,
|
||||
DAY_1_UNIX_SECS,
|
||||
),
|
||||
usage_user,
|
||||
sample_usage_row(
|
||||
"usage-user-b-fallback",
|
||||
"req-user-b-fallback",
|
||||
@@ -1272,6 +1431,61 @@ async fn gateway_handles_admin_stats_leaderboard_users_locally_without_user_read
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_stats_leaderboard_users_without_legacy_username_fallback_when_user_reader_exists(
|
||||
) {
|
||||
let (_upstream_url, upstream_hits, upstream_handle) =
|
||||
start_stats_upstream("/api/admin/stats/leaderboard/users").await;
|
||||
|
||||
let mut usage_user = sample_usage_row(
|
||||
"usage-user-no-legacy-fallback",
|
||||
"req-user-no-legacy-fallback",
|
||||
Some("user-missing"),
|
||||
Some("key-1"),
|
||||
Some("primary-key"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
60,
|
||||
20,
|
||||
0.4,
|
||||
0.4,
|
||||
DAY_1_UNIX_SECS,
|
||||
);
|
||||
usage_user.username = Some("stale-alice".to_string());
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage_user]));
|
||||
let data_state = GatewayDataState::with_usage_reader_for_tests(usage_repository)
|
||||
.with_user_reader(Arc::new(InMemoryUserReadRepository::seed_auth_users(
|
||||
Vec::<StoredUserAuthRecord>::new(),
|
||||
)));
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state)
|
||||
.without_auth_user_store_for_tests(),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(
|
||||
reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/stats/leaderboard/users?start_date=2024-03-21&end_date=2024-03-21&metric=cost&order=desc&tz_offset_minutes=0"
|
||||
)),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["items"][0]["id"], "user-missing");
|
||||
assert_eq!(payload["items"][0]["name"], "user-missing");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_stats_leaderboard_users_locally_without_usage_reader() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_crypto::{decrypt_python_fernet_ciphertext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::InMemoryAuthApiKeySnapshotRepository;
|
||||
use aether_data::repository::auth_modules::{
|
||||
AuthModuleReadRepository, InMemoryAuthModuleReadRepository, StoredOAuthProviderModuleConfig,
|
||||
};
|
||||
@@ -9,6 +10,8 @@ use aether_data::repository::oauth_providers::{
|
||||
InMemoryOAuthProviderRepository, OAuthProviderReadRepository, StoredOAuthProviderConfig,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::users::StoredUserAuthRecord;
|
||||
use aether_data::repository::wallet::{StoredWalletSnapshot, WalletLookupKey};
|
||||
use aether_data_contracts::repository::global_models::{
|
||||
AdminGlobalModelListQuery, AdminProviderModelListQuery, GlobalModelReadRepository,
|
||||
StoredPublicGlobalModel,
|
||||
@@ -213,6 +216,26 @@ fn fixture_system_import_payload(name: &str) -> Value {
|
||||
serde_json::from_str(raw).expect("fixture json should parse")
|
||||
}
|
||||
|
||||
fn sample_import_admin_user(user_id: &str) -> StoredUserAuthRecord {
|
||||
StoredUserAuthRecord::new(
|
||||
user_id.to_string(),
|
||||
Some("admin@example.com".to_string()),
|
||||
true,
|
||||
"admin".to_string(),
|
||||
Some("admin-hash".to_string()),
|
||||
"admin".to_string(),
|
||||
"local".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
false,
|
||||
Some(chrono::Utc::now()),
|
||||
Some(chrono::Utc::now()),
|
||||
)
|
||||
.expect("admin user should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_imports_admin_system_config_locally_and_persists_data() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
@@ -270,8 +293,9 @@ async fn gateway_imports_admin_system_config_locally_and_persists_data() {
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let status = response.status();
|
||||
let payload: Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(status, StatusCode::OK, "payload={payload}");
|
||||
assert_eq!(payload["message"], "配置导入成功");
|
||||
assert_eq!(payload["stats"]["global_models"]["created"], json!(1));
|
||||
assert_eq!(payload["stats"]["providers"]["created"], json!(1));
|
||||
@@ -504,6 +528,238 @@ async fn gateway_accepts_legacy_admin_system_config_import_versions() {
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_imports_admin_system_users_locally_and_persists_data() {
|
||||
let user_wallet_updated_at = "2024-05-06T07:08:09Z";
|
||||
let standalone_wallet_updated_at = "2024-06-07T08:09:10Z";
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().fallback(any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}));
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::default());
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_api_key_repository_for_tests(Arc::clone(&auth_repository))
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
)
|
||||
.with_auth_users_for_tests([sample_import_admin_user("admin-user-123")])
|
||||
.with_auth_wallets_for_tests(Vec::<StoredWalletSnapshot>::new());
|
||||
let gateway = build_router_with_state(state.clone());
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/admin/system/users/import"))
|
||||
.header(GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"version": "1.3",
|
||||
"merge_mode": "overwrite",
|
||||
"users": [{
|
||||
"email": "alice@example.com",
|
||||
"email_verified": true,
|
||||
"username": "alice",
|
||||
"password_hash": "argon2:imported-user-hash",
|
||||
"role": "user",
|
||||
"allowed_providers": ["openai"],
|
||||
"allowed_api_formats": ["openai:chat"],
|
||||
"allowed_models": ["gpt-5"],
|
||||
"rate_limit": 77,
|
||||
"is_active": true,
|
||||
"wallet": {
|
||||
"balance": 20.0,
|
||||
"recharge_balance": 15.0,
|
||||
"gift_balance": 5.0,
|
||||
"limit_mode": "finite",
|
||||
"currency": "CNY",
|
||||
"status": "locked",
|
||||
"total_recharged": 48.5,
|
||||
"total_consumed": 31.25,
|
||||
"total_refunded": 2.5,
|
||||
"total_adjusted": 7.75,
|
||||
"updated_at": user_wallet_updated_at
|
||||
},
|
||||
"api_keys": [{
|
||||
"key": "sk-user-import-1",
|
||||
"name": "Alice CLI",
|
||||
"allowed_providers": ["openai"],
|
||||
"allowed_api_formats": ["openai:chat"],
|
||||
"allowed_models": ["gpt-5"],
|
||||
"rate_limit": 60,
|
||||
"concurrent_limit": 3,
|
||||
"is_active": true,
|
||||
"expires_at": "2099-01-01T00:00:00Z",
|
||||
"auto_delete_on_expiry": false,
|
||||
"total_requests": 12,
|
||||
"total_cost_usd": 1.25
|
||||
}]
|
||||
}],
|
||||
"standalone_keys": [{
|
||||
"key": "sk-standalone-import-1",
|
||||
"name": "Imported Standalone",
|
||||
"allowed_providers": ["openai"],
|
||||
"allowed_api_formats": ["openai:chat"],
|
||||
"allowed_models": ["gpt-5"],
|
||||
"rate_limit": 90,
|
||||
"concurrent_limit": 4,
|
||||
"is_active": true,
|
||||
"expires_at": "2099-02-01T00:00:00Z",
|
||||
"auto_delete_on_expiry": false,
|
||||
"total_requests": 3,
|
||||
"total_cost_usd": 0.75,
|
||||
"wallet": {
|
||||
"balance": 30.0,
|
||||
"recharge_balance": 20.0,
|
||||
"gift_balance": 10.0,
|
||||
"limit_mode": "finite",
|
||||
"currency": "EUR",
|
||||
"status": "disabled",
|
||||
"total_recharged": 91.0,
|
||||
"total_consumed": 63.25,
|
||||
"total_refunded": 4.5,
|
||||
"total_adjusted": 13.0,
|
||||
"updated_at": standalone_wallet_updated_at
|
||||
}
|
||||
}]
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let payload: Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(status, StatusCode::OK, "payload={payload}");
|
||||
assert_eq!(payload["message"], "用户数据导入成功");
|
||||
assert_eq!(payload["stats"]["users"]["created"], json!(1));
|
||||
assert_eq!(payload["stats"]["api_keys"]["created"], json!(1));
|
||||
assert_eq!(payload["stats"]["standalone_keys"]["created"], json!(1));
|
||||
assert_eq!(payload["stats"]["errors"], json!([]));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let imported_user = state
|
||||
.find_user_auth_by_identifier("alice@example.com")
|
||||
.await
|
||||
.expect("user lookup should succeed")
|
||||
.expect("imported user should exist");
|
||||
assert_eq!(imported_user.username, "alice");
|
||||
assert_eq!(
|
||||
imported_user.password_hash.as_deref(),
|
||||
Some("argon2:imported-user-hash")
|
||||
);
|
||||
assert_eq!(imported_user.role, "user");
|
||||
assert_eq!(
|
||||
imported_user.allowed_providers,
|
||||
Some(vec!["openai".to_string()])
|
||||
);
|
||||
assert_eq!(
|
||||
imported_user.allowed_api_formats,
|
||||
Some(vec!["openai:chat".to_string()])
|
||||
);
|
||||
assert_eq!(
|
||||
imported_user.allowed_models,
|
||||
Some(vec!["gpt-5".to_string()])
|
||||
);
|
||||
assert!(imported_user.is_active);
|
||||
|
||||
let user_wallet = state
|
||||
.find_wallet(WalletLookupKey::UserId(&imported_user.id))
|
||||
.await
|
||||
.expect("user wallet lookup should succeed")
|
||||
.expect("user wallet should exist");
|
||||
assert_eq!(user_wallet.balance, 15.0);
|
||||
assert_eq!(user_wallet.gift_balance, 5.0);
|
||||
assert_eq!(user_wallet.limit_mode, "finite");
|
||||
assert_eq!(user_wallet.currency, "CNY");
|
||||
assert_eq!(user_wallet.status, "locked");
|
||||
assert_eq!(user_wallet.total_recharged, 48.5);
|
||||
assert_eq!(user_wallet.total_consumed, 31.25);
|
||||
assert_eq!(user_wallet.total_refunded, 2.5);
|
||||
assert_eq!(user_wallet.total_adjusted, 7.75);
|
||||
assert_eq!(
|
||||
user_wallet.updated_at_unix_secs,
|
||||
chrono::DateTime::parse_from_rfc3339(user_wallet_updated_at)
|
||||
.expect("user wallet updated_at should parse")
|
||||
.timestamp() as u64
|
||||
);
|
||||
|
||||
let user_api_keys = state
|
||||
.list_auth_api_key_export_records_by_user_ids(std::slice::from_ref(&imported_user.id))
|
||||
.await
|
||||
.expect("user api keys should load");
|
||||
assert_eq!(user_api_keys.len(), 1);
|
||||
assert_eq!(user_api_keys[0].name.as_deref(), Some("Alice CLI"));
|
||||
assert_eq!(
|
||||
user_api_keys[0].allowed_api_formats,
|
||||
Some(vec!["openai:chat".to_string()])
|
||||
);
|
||||
assert_eq!(
|
||||
decrypt_python_fernet_ciphertext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
user_api_keys[0]
|
||||
.key_encrypted
|
||||
.as_deref()
|
||||
.expect("encrypted user api key should exist"),
|
||||
)
|
||||
.expect("user api key should decrypt"),
|
||||
"sk-user-import-1"
|
||||
);
|
||||
|
||||
let standalone_keys = state
|
||||
.list_auth_api_key_export_standalone_records()
|
||||
.await
|
||||
.expect("standalone api keys should load");
|
||||
assert_eq!(standalone_keys.len(), 1);
|
||||
assert_eq!(
|
||||
standalone_keys[0].name.as_deref(),
|
||||
Some("Imported Standalone")
|
||||
);
|
||||
assert_eq!(
|
||||
decrypt_python_fernet_ciphertext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
standalone_keys[0]
|
||||
.key_encrypted
|
||||
.as_deref()
|
||||
.expect("encrypted standalone api key should exist"),
|
||||
)
|
||||
.expect("standalone api key should decrypt"),
|
||||
"sk-standalone-import-1"
|
||||
);
|
||||
|
||||
let standalone_wallet = state
|
||||
.find_wallet(WalletLookupKey::ApiKeyId(&standalone_keys[0].api_key_id))
|
||||
.await
|
||||
.expect("standalone wallet lookup should succeed")
|
||||
.expect("standalone wallet should exist");
|
||||
assert_eq!(standalone_wallet.balance, 20.0);
|
||||
assert_eq!(standalone_wallet.gift_balance, 10.0);
|
||||
assert_eq!(standalone_wallet.limit_mode, "finite");
|
||||
assert_eq!(standalone_wallet.currency, "EUR");
|
||||
assert_eq!(standalone_wallet.status, "disabled");
|
||||
assert_eq!(standalone_wallet.total_recharged, 91.0);
|
||||
assert_eq!(standalone_wallet.total_consumed, 63.25);
|
||||
assert_eq!(standalone_wallet.total_refunded, 4.5);
|
||||
assert_eq!(standalone_wallet.total_adjusted, 13.0);
|
||||
assert_eq!(
|
||||
standalone_wallet.updated_at_unix_secs,
|
||||
chrono::DateTime::parse_from_rfc3339(standalone_wallet_updated_at)
|
||||
.expect("standalone wallet updated_at should parse")
|
||||
.timestamp() as u64
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
let _ = upstream_url;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_imports_admin_system_config_fixtures_from_legacy_exports() {
|
||||
for fixture in ["v20", "v21", "v22"] {
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||
use aether_data::repository::users::{InMemoryUserReadRepository, StoredUserSummary};
|
||||
use aether_data::repository::users::{
|
||||
InMemoryUserReadRepository, StoredUserAuthRecord, StoredUserSummary,
|
||||
};
|
||||
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::routing::{any, get, post};
|
||||
@@ -208,6 +213,57 @@ fn sample_user_summary(id: &str, username: &str) -> StoredUserSummary {
|
||||
.expect("user summary should build")
|
||||
}
|
||||
|
||||
fn sample_auth_user(id: &str, username: &str, role: &str, is_active: bool) -> StoredUserAuthRecord {
|
||||
StoredUserAuthRecord::new(
|
||||
id.to_string(),
|
||||
Some(format!("{username}@example.com")),
|
||||
true,
|
||||
username.to_string(),
|
||||
Some("hash".to_string()),
|
||||
role.to_string(),
|
||||
"local".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
is_active,
|
||||
false,
|
||||
Some(chrono::Utc::now()),
|
||||
Some(chrono::Utc::now()),
|
||||
)
|
||||
.expect("auth user should build")
|
||||
}
|
||||
|
||||
fn sample_auth_api_key_snapshot(
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
api_key_name: &str,
|
||||
) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
format!("snapshot-{user_id}"),
|
||||
Some(format!("snapshot-{user_id}@example.com")),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
api_key_id.to_string(),
|
||||
Some(api_key_name.to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn recent_unix_secs(minutes_ago: u64) -> i64 {
|
||||
let now = chrono::Utc::now().timestamp();
|
||||
now.saturating_sub((minutes_ago * 60) as i64)
|
||||
@@ -666,13 +722,18 @@ async fn gateway_handles_admin_usage_active_locally_with_trusted_admin_principal
|
||||
)],
|
||||
vec![provider_key],
|
||||
));
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-key-1".to_string()),
|
||||
sample_auth_api_key_snapshot("user-1", "key-1", "fresh-primary"),
|
||||
)]));
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_usage_reader_for_tests(usage_repository)
|
||||
.with_provider_catalog_reader(provider_catalog_repository),
|
||||
.with_provider_catalog_reader(provider_catalog_repository)
|
||||
.with_auth_api_key_reader(auth_repository),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
@@ -689,6 +750,7 @@ async fn gateway_handles_admin_usage_active_locally_with_trusted_admin_principal
|
||||
assert_eq!(payload["requests"][0]["id"], "usage-pending");
|
||||
assert_eq!(payload["requests"][0]["effective_input_tokens"], 5);
|
||||
assert_eq!(payload["requests"][0]["provider"], "OpenAI");
|
||||
assert_eq!(payload["requests"][0]["api_key_name"], "fresh-primary");
|
||||
assert_eq!(
|
||||
payload["requests"][0]["provider_key_name"],
|
||||
"upstream-primary"
|
||||
@@ -847,6 +909,120 @@ async fn gateway_handles_admin_usage_records_with_provider_key_name_fallback_fro
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_usage_records_with_snapshot_first_user_and_api_key_names() {
|
||||
let (_upstream_url, upstream_hits, upstream_handle) =
|
||||
start_usage_upstream("/api/admin/usage/records").await;
|
||||
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![sample_usage_row(
|
||||
"usage-snapshot",
|
||||
"req-snapshot",
|
||||
Some("user-2"),
|
||||
Some("key-2"),
|
||||
Some("legacy-secondary"),
|
||||
"Anthropic",
|
||||
"claude-3-7",
|
||||
"failed",
|
||||
40,
|
||||
10,
|
||||
0.1,
|
||||
0.12,
|
||||
DAY_2_UNIX_SECS,
|
||||
)]));
|
||||
let user_repository = Arc::new(InMemoryUserReadRepository::seed(vec![sample_user_summary(
|
||||
"user-2",
|
||||
"fresh-bob",
|
||||
)]));
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-key-2".to_string()),
|
||||
sample_auth_api_key_snapshot("user-2", "key-2", "fresh-secondary"),
|
||||
)]));
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_usage_reader_for_tests(usage_repository)
|
||||
.with_user_reader(user_repository)
|
||||
.with_auth_api_key_reader(auth_repository),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/usage/records?search=fresh-secondary&username=fresh-bob&limit=10&offset=0"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["records"][0]["id"], "usage-snapshot");
|
||||
assert_eq!(payload["records"][0]["username"], "fresh-bob");
|
||||
assert_eq!(payload["records"][0]["api_key"]["name"], "fresh-secondary");
|
||||
assert_eq!(payload["records"][0]["api_key_name"], "fresh-secondary");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_usage_records_without_legacy_username_fallback_when_user_reader_exists(
|
||||
) {
|
||||
let (_upstream_url, upstream_hits, upstream_handle) =
|
||||
start_usage_upstream("/api/admin/usage/records").await;
|
||||
|
||||
let mut usage = sample_usage_row(
|
||||
"usage-no-legacy-username",
|
||||
"req-no-legacy-username",
|
||||
Some("user-missing"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
10,
|
||||
2,
|
||||
0.01,
|
||||
0.012,
|
||||
DAY_1_UNIX_SECS,
|
||||
);
|
||||
usage.username = Some("stale-alice".to_string());
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_usage_reader_for_tests(Arc::new(
|
||||
InMemoryUsageReadRepository::seed(vec![usage]),
|
||||
))
|
||||
.with_user_reader(Arc::new(
|
||||
InMemoryUserReadRepository::seed_auth_users(Vec::<StoredUserAuthRecord>::new()),
|
||||
)),
|
||||
)
|
||||
.without_auth_user_store_for_tests(),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/usage/records?username=stale-alice&limit=10&offset=0"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], 0);
|
||||
assert!(payload["records"].as_array().expect("array").is_empty());
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_usage_detail_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
@@ -912,6 +1088,28 @@ async fn gateway_handles_admin_usage_detail_locally_with_trusted_admin_principal
|
||||
}));
|
||||
usage.request_metadata = Some(json!({
|
||||
"trace_id": "trace-123",
|
||||
"candidate_id": "cand-detail-1",
|
||||
"candidate_index": 2,
|
||||
"key_name": "upstream-primary",
|
||||
"planner_kind": "claude_cli_sync",
|
||||
"route_family": "claude",
|
||||
"route_kind": "cli",
|
||||
"execution_path": "local_execution_runtime_miss",
|
||||
"local_execution_runtime_miss_reason": "all_candidates_skipped",
|
||||
"billing_snapshot_schema_version": "v2",
|
||||
"billing_snapshot_status": "resolved",
|
||||
"rate_multiplier": 0.5,
|
||||
"is_free_tier": false,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 15.0,
|
||||
"cache_creation_price_per_1m": 3.75,
|
||||
"cache_read_price_per_1m": 0.3,
|
||||
"price_per_request": 0.02,
|
||||
"billing_snapshot": {
|
||||
"resolved_variables": {
|
||||
"output_price_per_1m": 11.0
|
||||
}
|
||||
}
|
||||
}));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
|
||||
let user_repository = Arc::new(InMemoryUserReadRepository::seed(vec![sample_user_summary(
|
||||
@@ -962,16 +1160,74 @@ async fn gateway_handles_admin_usage_detail_locally_with_trusted_admin_principal
|
||||
);
|
||||
assert_eq!(payload["response_headers"]["X-Upstream"], "openai");
|
||||
assert_eq!(payload["client_response_headers"]["X-Request-Id"], "req-1");
|
||||
assert!(payload["metadata"]["trace_id"].is_null());
|
||||
assert_eq!(payload["trace"]["trace_id"], "trace-123");
|
||||
assert!(payload["metadata"]["request_preview_source"].is_null());
|
||||
assert!(payload["metadata"]["original_request_body_available"].is_null());
|
||||
assert!(payload["metadata"]["original_response_body_available"].is_null());
|
||||
assert!(payload["metadata"]["candidate_id"].is_null());
|
||||
assert!(payload["metadata"]["candidate_index"].is_null());
|
||||
assert!(payload["metadata"]["key_name"].is_null());
|
||||
assert!(payload["metadata"]["planner_kind"].is_null());
|
||||
assert!(payload["metadata"]["route_family"].is_null());
|
||||
assert!(payload["metadata"]["route_kind"].is_null());
|
||||
assert!(payload["metadata"]["execution_path"].is_null());
|
||||
assert!(payload["metadata"]["local_execution_runtime_miss_reason"].is_null());
|
||||
assert!(payload["metadata"]["billing_snapshot"].is_null());
|
||||
assert!(payload["metadata"]["billing_snapshot_schema_version"].is_null());
|
||||
assert!(payload["metadata"]["billing_snapshot_status"].is_null());
|
||||
assert!(payload["metadata"]["rate_multiplier"].is_null());
|
||||
assert!(payload["metadata"]["is_free_tier"].is_null());
|
||||
assert!(payload["metadata"]["input_price_per_1m"].is_null());
|
||||
assert!(payload["metadata"]["output_price_per_1m"].is_null());
|
||||
assert!(payload["metadata"]["cache_creation_price_per_1m"].is_null());
|
||||
assert!(payload["metadata"]["cache_read_price_per_1m"].is_null());
|
||||
assert!(payload["metadata"]["price_per_request"].is_null());
|
||||
assert_eq!(payload["body_capture"]["request"]["storage"], "inline");
|
||||
assert_eq!(payload["body_capture"]["request"]["available"], true);
|
||||
assert_eq!(
|
||||
payload["metadata"]["request_preview_source"],
|
||||
payload["body_capture"]["request"]["preview_source"],
|
||||
"stored_original"
|
||||
);
|
||||
assert_eq!(payload["metadata"]["trace_id"], "trace-123");
|
||||
assert_eq!(payload["metadata"]["original_request_body_available"], true);
|
||||
assert_eq!(
|
||||
payload["metadata"]["original_response_body_available"],
|
||||
true
|
||||
payload["body_capture"]["provider_request"]["storage"],
|
||||
"inline"
|
||||
);
|
||||
assert_eq!(payload["body_capture"]["response"]["storage"], "inline");
|
||||
assert_eq!(
|
||||
payload["body_capture"]["client_response"]["storage"],
|
||||
"inline"
|
||||
);
|
||||
assert_eq!(payload["routing"]["candidate_id"], "cand-detail-1");
|
||||
assert_eq!(payload["routing"]["candidate_index"], 2);
|
||||
assert_eq!(payload["routing"]["key_name"], "upstream-primary");
|
||||
assert_eq!(payload["routing"]["planner_kind"], "claude_cli_sync");
|
||||
assert_eq!(payload["routing"]["route_family"], "claude");
|
||||
assert_eq!(payload["routing"]["route_kind"], "cli");
|
||||
assert_eq!(
|
||||
payload["routing"]["execution_path"],
|
||||
"local_execution_runtime_miss"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["routing"]["local_execution_runtime_miss_reason"],
|
||||
"all_candidates_skipped"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["settlement"]["billing_snapshot_schema_version"],
|
||||
"v2"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["settlement"]["billing_snapshot"]["resolved_variables"]["output_price_per_1m"],
|
||||
11.0
|
||||
);
|
||||
assert_eq!(payload["settlement"]["billing_snapshot_status"], "resolved");
|
||||
assert_eq!(payload["settlement"]["rate_multiplier"], 0.5);
|
||||
assert_eq!(payload["settlement"]["is_free_tier"], false);
|
||||
assert_eq!(payload["settlement"]["input_price_per_1m"], 3.0);
|
||||
assert_eq!(payload["settlement"]["output_price_per_1m"], 15.0);
|
||||
assert_eq!(payload["settlement"]["cache_creation_price_per_1m"], 3.75);
|
||||
assert_eq!(payload["settlement"]["cache_read_price_per_1m"], 0.3);
|
||||
assert_eq!(payload["settlement"]["price_per_request"], 0.02);
|
||||
assert_eq!(payload["has_request_body"], true);
|
||||
assert_eq!(payload["has_provider_request_body"], true);
|
||||
assert_eq!(payload["has_response_body"], true);
|
||||
@@ -1008,6 +1264,124 @@ async fn gateway_handles_admin_usage_detail_locally_with_trusted_admin_principal
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_usage_detail_with_ref_backed_bodies() {
|
||||
let (_upstream_url, upstream_hits, upstream_handle) =
|
||||
start_usage_upstream("/api/admin/usage/usage-ref-detail").await;
|
||||
|
||||
let mut usage = sample_usage_row(
|
||||
"usage-ref-detail",
|
||||
"req-ref-detail",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
120,
|
||||
30,
|
||||
0.3,
|
||||
0.36,
|
||||
DAY_1_UNIX_SECS,
|
||||
);
|
||||
usage.request_body = Some(json!({
|
||||
"model": "gpt-5",
|
||||
"messages": [{"role": "user", "content": "ref-backed request body"}],
|
||||
"stream": false,
|
||||
}));
|
||||
usage.provider_request_body = Some(json!({
|
||||
"model": "gpt-5-target",
|
||||
"temperature": 0.2,
|
||||
"stream": false,
|
||||
}));
|
||||
usage.response_body = Some(json!({
|
||||
"id": "resp-ref-detail",
|
||||
"usage": {"total_tokens": 150},
|
||||
}));
|
||||
usage.client_response_body = Some(json!({
|
||||
"id": "resp-ref-detail",
|
||||
"output_text": "hello from ref",
|
||||
}));
|
||||
usage.request_metadata = Some(json!({
|
||||
"trace_id": "trace-ref-detail",
|
||||
}));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed_with_detached_bodies(
|
||||
vec![usage],
|
||||
));
|
||||
let user_repository = Arc::new(InMemoryUserReadRepository::seed(vec![sample_user_summary(
|
||||
"user-1", "alice",
|
||||
)]));
|
||||
let data_state = GatewayDataState::with_usage_reader_for_tests(usage_repository)
|
||||
.with_user_reader(user_repository);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(
|
||||
reqwest::Client::new().get(format!("{gateway_url}/api/admin/usage/usage-ref-detail")),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert!(payload["metadata"]["trace_id"].is_null());
|
||||
assert_eq!(payload["trace"]["trace_id"], "trace-ref-detail");
|
||||
assert!(payload["metadata"]["request_preview_source"].is_null());
|
||||
assert!(payload["metadata"]["original_request_body_available"].is_null());
|
||||
assert!(payload["metadata"]["original_response_body_available"].is_null());
|
||||
assert_eq!(payload["body_capture"]["request"]["storage"], "reference");
|
||||
assert_eq!(payload["body_capture"]["request"]["available"], true);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["request"]["preview_source"],
|
||||
"stored_reference"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["provider_request"]["storage"],
|
||||
"reference"
|
||||
);
|
||||
assert_eq!(payload["body_capture"]["response"]["storage"], "reference");
|
||||
assert_eq!(
|
||||
payload["body_capture"]["client_response"]["storage"],
|
||||
"reference"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["request"]["body_ref"],
|
||||
"usage://request/req-ref-detail/request_body"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["provider_request"]["body_ref"],
|
||||
"usage://request/req-ref-detail/provider_request_body"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["response"]["body_ref"],
|
||||
"usage://request/req-ref-detail/response_body"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["client_response"]["body_ref"],
|
||||
"usage://request/req-ref-detail/client_response_body"
|
||||
);
|
||||
assert_eq!(payload["request_body"]["model"], "gpt-5");
|
||||
assert_eq!(
|
||||
payload["request_body"]["messages"][0]["content"],
|
||||
"ref-backed request body"
|
||||
);
|
||||
assert_eq!(payload["provider_request_body"]["temperature"], 0.2);
|
||||
assert_eq!(payload["response_body"]["id"], "resp-ref-detail");
|
||||
assert_eq!(
|
||||
payload["client_response_body"]["output_text"],
|
||||
"hello from ref"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_admin_usage_detail_attaches_explicit_audit() {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![sample_usage_row(
|
||||
@@ -1170,6 +1544,107 @@ async fn gateway_handles_admin_usage_replay_locally_with_trusted_admin_principal
|
||||
assert_eq!(payload["request_body"]["api_format"], "openai:chat");
|
||||
assert_eq!(payload["request_body"]["stream"], false);
|
||||
assert_eq!(payload["original_request_body_available"], false);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["request"]["preview_source"],
|
||||
"local_reconstruction"
|
||||
);
|
||||
assert_eq!(payload["body_capture"]["request"]["storage"], "missing");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_usage_replay_with_ref_backed_request_body() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/usage/usage-ref-replay/replay",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut usage = sample_usage_row(
|
||||
"usage-ref-replay",
|
||||
"req-ref-replay",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"success",
|
||||
120,
|
||||
50,
|
||||
0.42,
|
||||
0.37,
|
||||
DAY_1_UNIX_SECS,
|
||||
);
|
||||
usage.request_body = Some(json!({
|
||||
"model": "gpt-5",
|
||||
"messages": [{"role": "user", "content": "replay from ref"}],
|
||||
"stream": false,
|
||||
}));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed_with_detached_bodies(
|
||||
vec![usage],
|
||||
));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1", "OpenAI", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-1",
|
||||
"provider-1",
|
||||
"openai:chat",
|
||||
"https://api.openai.com/v1",
|
||||
)],
|
||||
vec![],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_usage_reader_for_tests(usage_repository)
|
||||
.with_provider_catalog_reader(provider_catalog_repository),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(
|
||||
reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/usage/usage-ref-replay/replay"
|
||||
))
|
||||
.json(&json!({
|
||||
"provider_id": "provider-1"
|
||||
})),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["request_body"]["model"], "gpt-5");
|
||||
assert_eq!(
|
||||
payload["request_body"]["messages"][0]["content"],
|
||||
"replay from ref"
|
||||
);
|
||||
assert_eq!(payload["original_request_body_available"], true);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["request"]["preview_source"],
|
||||
"stored_reference"
|
||||
);
|
||||
assert_eq!(payload["body_capture"]["request"]["storage"], "reference");
|
||||
assert_eq!(
|
||||
payload["body_capture"]["request"]["body_ref"],
|
||||
"usage://request/req-ref-replay/request_body"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
@@ -1315,6 +1790,16 @@ async fn gateway_handles_admin_usage_curl_locally_with_trusted_admin_principal()
|
||||
assert_eq!(payload["body"]["temperature"], 0.2);
|
||||
assert_eq!(payload["body"]["stream"], false);
|
||||
assert_eq!(payload["original_request_body_available"], true);
|
||||
assert_eq!(payload["body_capture"]["body_source"], "provider_request");
|
||||
assert_eq!(payload["body_capture"]["request"]["storage"], "inline");
|
||||
assert_eq!(
|
||||
payload["body_capture"]["request"]["preview_source"],
|
||||
"stored_original"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["provider_request"]["storage"],
|
||||
"inline"
|
||||
);
|
||||
let curl = payload["curl"].as_str().expect("curl should be string");
|
||||
assert!(curl.contains("curl"));
|
||||
assert!(curl.contains("https://api.openai.example/v1/chat/completions"));
|
||||
@@ -1329,6 +1814,103 @@ async fn gateway_handles_admin_usage_curl_locally_with_trusted_admin_principal()
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_usage_curl_with_ref_backed_provider_request_body() {
|
||||
let (_upstream_url, upstream_hits, upstream_handle) =
|
||||
start_usage_upstream("/api/admin/usage/usage-ref-curl/curl").await;
|
||||
|
||||
let mut usage = sample_usage_row(
|
||||
"usage-ref-curl",
|
||||
"req-ref-curl",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
120,
|
||||
30,
|
||||
0.3,
|
||||
0.36,
|
||||
DAY_1_UNIX_SECS,
|
||||
);
|
||||
usage.request_body = Some(json!({
|
||||
"model": "gpt-5",
|
||||
"messages": [{"role": "user", "content": "client body"}],
|
||||
"stream": false,
|
||||
}));
|
||||
usage.provider_request_headers = Some(json!({
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": "Bearer provider-real",
|
||||
}));
|
||||
usage.provider_request_body = Some(json!({
|
||||
"model": "gpt-5-target",
|
||||
"stream": false,
|
||||
"temperature": 0.2,
|
||||
}));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed_with_detached_bodies(
|
||||
vec![usage],
|
||||
));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1", "openai", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-1",
|
||||
"provider-1",
|
||||
"openai:chat",
|
||||
"https://api.openai.example",
|
||||
)],
|
||||
vec![sample_key(
|
||||
"provider-key-1",
|
||||
"provider-1",
|
||||
"openai:chat",
|
||||
"sk-test",
|
||||
)],
|
||||
));
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_and_usage_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
usage_repository,
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(
|
||||
reqwest::Client::new().get(format!("{gateway_url}/api/admin/usage/usage-ref-curl/curl")),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["body"]["model"], "gpt-5-target");
|
||||
assert_eq!(payload["body"]["temperature"], 0.2);
|
||||
assert_eq!(payload["original_request_body_available"], true);
|
||||
assert_eq!(payload["body_capture"]["body_source"], "provider_request");
|
||||
assert_eq!(payload["body_capture"]["request"]["storage"], "reference");
|
||||
assert_eq!(
|
||||
payload["body_capture"]["request"]["preview_source"],
|
||||
"stored_reference"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["provider_request"]["storage"],
|
||||
"reference"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["body_capture"]["provider_request"]["body_ref"],
|
||||
"usage://request/req-ref-curl/provider_request_body"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_admin_usage_curl_attaches_explicit_audit() {
|
||||
let mut usage = sample_usage_row(
|
||||
@@ -1616,6 +2198,147 @@ async fn gateway_handles_admin_usage_cache_affinity_interval_timeline_locally_wi
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_usage_cache_affinity_interval_timeline_with_auth_user_fallback() {
|
||||
let (_upstream_url, upstream_hits, upstream_handle) =
|
||||
start_usage_upstream("/api/admin/usage/cache-affinity/interval-timeline").await;
|
||||
let mut usage_one = sample_usage_row(
|
||||
"usage-auth-fallback-1",
|
||||
"req-auth-fallback-1",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
10,
|
||||
2,
|
||||
0.01,
|
||||
0.012,
|
||||
recent_unix_secs(55),
|
||||
);
|
||||
usage_one.username = Some("stale-alice".to_string());
|
||||
let mut usage_two = sample_usage_row(
|
||||
"usage-auth-fallback-2",
|
||||
"req-auth-fallback-2",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
12,
|
||||
3,
|
||||
0.01,
|
||||
0.012,
|
||||
recent_unix_secs(50),
|
||||
);
|
||||
usage_two.username = Some("stale-alice".to_string());
|
||||
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
usage_one, usage_two,
|
||||
]));
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(
|
||||
usage_repository,
|
||||
))
|
||||
.with_auth_users_for_tests([sample_auth_user("user-1", "fresh-alice", "user", true)]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/usage/cache-affinity/interval-timeline?hours=24&limit=100&include_user_info=true"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total_points"], 1);
|
||||
assert_eq!(payload["points"][0]["user_id"], "user-1");
|
||||
assert_eq!(payload["users"]["user-1"], "fresh-alice");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_usage_cache_affinity_interval_timeline_without_legacy_username_fallback_when_user_reader_exists(
|
||||
) {
|
||||
let (_upstream_url, upstream_hits, upstream_handle) =
|
||||
start_usage_upstream("/api/admin/usage/cache-affinity/interval-timeline").await;
|
||||
let mut usage_one = sample_usage_row(
|
||||
"usage-auth-reader-miss-1",
|
||||
"req-auth-reader-miss-1",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
10,
|
||||
2,
|
||||
0.01,
|
||||
0.012,
|
||||
recent_unix_secs(55),
|
||||
);
|
||||
usage_one.username = Some("stale-alice".to_string());
|
||||
let mut usage_two = sample_usage_row(
|
||||
"usage-auth-reader-miss-2",
|
||||
"req-auth-reader-miss-2",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
12,
|
||||
3,
|
||||
0.01,
|
||||
0.012,
|
||||
recent_unix_secs(50),
|
||||
);
|
||||
usage_two.username = Some("stale-alice".to_string());
|
||||
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
usage_one, usage_two,
|
||||
]));
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(
|
||||
usage_repository,
|
||||
))
|
||||
.with_auth_users_for_tests([sample_auth_user("user-2", "fresh-bob", "user", true)]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/usage/cache-affinity/interval-timeline?hours=24&limit=100&include_user_info=true"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total_points"], 1);
|
||||
assert_eq!(payload["points"][0]["user_id"], "user-1");
|
||||
assert!(payload["users"]
|
||||
.as_object()
|
||||
.expect("users should be object")
|
||||
.get("user-1")
|
||||
.is_none());
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_usage_cache_affinity_ttl_analysis_locally_with_trusted_admin_principal(
|
||||
) {
|
||||
|
||||
@@ -2156,6 +2156,37 @@ fn sample_auth_session(
|
||||
.expect("auth session should build")
|
||||
}
|
||||
|
||||
fn sample_usage_auth_snapshot(
|
||||
api_key_id: &str,
|
||||
user_id: &str,
|
||||
api_key_name: &str,
|
||||
) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(json!(["openai"])),
|
||||
Some(json!(["openai:chat"])),
|
||||
Some(json!(["gpt-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some(api_key_name.to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
None,
|
||||
Some(json!(["openai"])),
|
||||
Some(json!(["openai:chat"])),
|
||||
Some(json!(["gpt-5"])),
|
||||
)
|
||||
.expect("auth api key snapshot should build")
|
||||
}
|
||||
|
||||
async fn start_auth_gateway_with_state(
|
||||
user: StoredUserAuthRecord,
|
||||
wallet: StoredWalletSnapshot,
|
||||
@@ -4671,6 +4702,22 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() {
|
||||
]),
|
||||
now + chrono::Duration::hours(1),
|
||||
);
|
||||
let mut streaming_usage = sample_user_usage_audit(
|
||||
"usage-users-me-streaming-1",
|
||||
"req-users-me-streaming-1",
|
||||
"user-auth-1",
|
||||
"gpt-4.1-mini",
|
||||
"OpenAI",
|
||||
"streaming",
|
||||
now - chrono::Duration::minutes(5),
|
||||
);
|
||||
streaming_usage.request_metadata = Some(json!({
|
||||
"rate_multiplier": 0.5,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 9.0,
|
||||
"cache_creation_price_per_1m": 3.75,
|
||||
"cache_read_price_per_1m": 0.3,
|
||||
}));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_user_usage_audit(
|
||||
"usage-users-me-completed-1",
|
||||
@@ -4690,34 +4737,41 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() {
|
||||
"failed",
|
||||
now - chrono::Duration::minutes(10),
|
||||
),
|
||||
sample_user_usage_audit(
|
||||
"usage-users-me-streaming-1",
|
||||
"req-users-me-streaming-1",
|
||||
"user-auth-1",
|
||||
"gpt-4.1-mini",
|
||||
"OpenAI",
|
||||
"streaming",
|
||||
now - chrono::Duration::minutes(5),
|
||||
),
|
||||
streaming_usage,
|
||||
]));
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-api-key-user-1".to_string()),
|
||||
sample_usage_auth_snapshot("api-key-user-1", "user-auth-1", "renamed-key"),
|
||||
)]));
|
||||
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
|
||||
start_auth_gateway_with_usage_state(
|
||||
user,
|
||||
sample_auth_wallet("user-auth-1", now),
|
||||
[sample_auth_session(
|
||||
"user-auth-1",
|
||||
"session-users-me-usage-1",
|
||||
"device-users-me-usage-1",
|
||||
"refresh-token-placeholder",
|
||||
now,
|
||||
)],
|
||||
usage_repository,
|
||||
)
|
||||
start_auth_gateway_with_builder(|| {
|
||||
let data_state = crate::data::GatewayDataState::with_user_wallet_and_usage_for_tests(
|
||||
Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
|
||||
user.clone()
|
||||
])),
|
||||
Arc::new(InMemoryWalletRepository::seed(vec![sample_auth_wallet(
|
||||
"user-auth-1",
|
||||
now,
|
||||
)])),
|
||||
Arc::clone(&usage_repository),
|
||||
)
|
||||
.with_auth_api_key_reader(auth_repository);
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_auth_sessions_for_tests([sample_auth_session(
|
||||
"user-auth-1",
|
||||
"session-users-me-usage-1",
|
||||
"device-users-me-usage-1",
|
||||
"refresh-token-placeholder",
|
||||
now,
|
||||
)])
|
||||
})
|
||||
.await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/users/me/usage?limit=10&offset=0"
|
||||
"{gateway_url}/api/users/me/usage?limit=10&offset=0&search=renamed-key"
|
||||
))
|
||||
.header("authorization", format!("Bearer {access_token}"))
|
||||
.header("x-client-device-id", "device-users-me-usage-1")
|
||||
@@ -4744,6 +4798,12 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() {
|
||||
payload["records"][0]["cache_creation_ephemeral_1h_input_tokens"],
|
||||
6
|
||||
);
|
||||
assert_eq!(payload["records"][0]["input_price_per_1m"], 3.0);
|
||||
assert_eq!(payload["records"][0]["output_price_per_1m"], 9.0);
|
||||
assert_eq!(payload["records"][0]["cache_creation_price_per_1m"], 3.75);
|
||||
assert_eq!(payload["records"][0]["cache_read_price_per_1m"], 0.3);
|
||||
assert_eq!(payload["records"][0]["api_key"]["name"], "renamed-key");
|
||||
assert_eq!(payload["records"][0]["api_key"]["display"], "renamed-key");
|
||||
assert_eq!(
|
||||
payload["summary_by_model"]
|
||||
.as_array()
|
||||
@@ -4771,6 +4831,91 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() {
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_users_me_usage_without_legacy_api_key_name_fallback_when_auth_reader_exists(
|
||||
) {
|
||||
let now = Utc::now();
|
||||
let user = sample_auth_user(now);
|
||||
let access_token = build_test_auth_token(
|
||||
"access",
|
||||
serde_json::Map::from_iter([
|
||||
("user_id".to_string(), json!(user.id)),
|
||||
("role".to_string(), json!(user.role)),
|
||||
(
|
||||
"created_at".to_string(),
|
||||
json!(user.created_at.map(|value| value.to_rfc3339())),
|
||||
),
|
||||
(
|
||||
"session_id".to_string(),
|
||||
json!("session-users-me-usage-no-legacy-fallback"),
|
||||
),
|
||||
]),
|
||||
now + chrono::Duration::hours(1),
|
||||
);
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_user_usage_audit(
|
||||
"usage-users-me-completed-no-legacy-fallback",
|
||||
"req-users-me-completed-no-legacy-fallback",
|
||||
"user-auth-1",
|
||||
"gpt-4.1",
|
||||
"OpenAI",
|
||||
"completed",
|
||||
now - chrono::Duration::minutes(20),
|
||||
),
|
||||
]));
|
||||
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
|
||||
start_auth_gateway_with_builder(|| {
|
||||
let data_state = crate::data::GatewayDataState::with_user_wallet_and_usage_for_tests(
|
||||
Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
|
||||
user.clone()
|
||||
])),
|
||||
Arc::new(InMemoryWalletRepository::seed(vec![sample_auth_wallet(
|
||||
"user-auth-1",
|
||||
now,
|
||||
)])),
|
||||
Arc::clone(&usage_repository),
|
||||
)
|
||||
.with_auth_api_key_reader(Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![])));
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_auth_sessions_for_tests([sample_auth_session(
|
||||
"user-auth-1",
|
||||
"session-users-me-usage-no-legacy-fallback",
|
||||
"device-users-me-usage-no-legacy-fallback",
|
||||
"refresh-token-users-me-usage-no-legacy-fallback",
|
||||
now,
|
||||
)])
|
||||
})
|
||||
.await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/users/me/usage?limit=10&offset=0&search=default"
|
||||
))
|
||||
.header("authorization", format!("Bearer {access_token}"))
|
||||
.header(
|
||||
"x-client-device-id",
|
||||
"device-users-me-usage-no-legacy-fallback",
|
||||
)
|
||||
.header("user-agent", "AetherTest/1.0")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["pagination"]["total"], 0);
|
||||
assert_eq!(
|
||||
payload["records"].as_array().expect("records array").len(),
|
||||
0
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_users_me_usage_active_locally_without_proxying_upstream() {
|
||||
let now = Utc::now();
|
||||
@@ -4788,6 +4933,18 @@ async fn gateway_handles_users_me_usage_active_locally_without_proxying_upstream
|
||||
]),
|
||||
now + chrono::Duration::hours(1),
|
||||
);
|
||||
let mut streaming_usage = sample_user_usage_audit(
|
||||
"usage-users-me-streaming-1",
|
||||
"req-users-me-streaming-1",
|
||||
"user-auth-1",
|
||||
"gpt-4.1-mini",
|
||||
"OpenAI",
|
||||
"streaming",
|
||||
now - chrono::Duration::minutes(2),
|
||||
);
|
||||
streaming_usage.request_metadata = Some(json!({
|
||||
"rate_multiplier": 0.5,
|
||||
}));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_user_usage_audit(
|
||||
"usage-users-me-pending-1",
|
||||
@@ -4798,15 +4955,7 @@ async fn gateway_handles_users_me_usage_active_locally_without_proxying_upstream
|
||||
"pending",
|
||||
now - chrono::Duration::minutes(4),
|
||||
),
|
||||
sample_user_usage_audit(
|
||||
"usage-users-me-streaming-1",
|
||||
"req-users-me-streaming-1",
|
||||
"user-auth-1",
|
||||
"gpt-4.1-mini",
|
||||
"OpenAI",
|
||||
"streaming",
|
||||
now - chrono::Duration::minutes(2),
|
||||
),
|
||||
streaming_usage,
|
||||
sample_user_usage_audit(
|
||||
"usage-users-me-completed-1",
|
||||
"req-users-me-completed-1",
|
||||
@@ -4846,6 +4995,7 @@ async fn gateway_handles_users_me_usage_active_locally_without_proxying_upstream
|
||||
let requests = payload["requests"].as_array().expect("requests array");
|
||||
assert_eq!(requests.len(), 2);
|
||||
assert_eq!(requests[0]["status"], "streaming");
|
||||
assert_eq!(requests[0]["rate_multiplier"], 0.5);
|
||||
assert_eq!(requests[0]["cache_creation_ephemeral_5m_input_tokens"], 4);
|
||||
assert_eq!(requests[0]["cache_creation_ephemeral_1h_input_tokens"], 6);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -641,6 +641,94 @@ async fn gateway_handles_dashboard_recent_requests_locally_without_proxying_upst
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_dashboard_recent_requests_with_auth_user_fallback() {
|
||||
let now = stable_dashboard_now();
|
||||
let user = sample_auth_user(now);
|
||||
let access_token = build_test_auth_token(
|
||||
"access",
|
||||
serde_json::Map::from_iter([
|
||||
("user_id".to_string(), json!(user.id)),
|
||||
("role".to_string(), json!(user.role)),
|
||||
(
|
||||
"created_at".to_string(),
|
||||
json!(user.created_at.map(|value| value.to_rfc3339())),
|
||||
),
|
||||
(
|
||||
"session_id".to_string(),
|
||||
json!("session-dashboard-recent-auth-fallback"),
|
||||
),
|
||||
]),
|
||||
chrono::Utc::now() + chrono::Duration::hours(1),
|
||||
);
|
||||
let session = sample_auth_session(
|
||||
"user-auth-1",
|
||||
"session-dashboard-recent-auth-fallback",
|
||||
"device-dashboard-recent-auth-fallback",
|
||||
"refresh-dashboard-recent-auth-fallback",
|
||||
now,
|
||||
);
|
||||
let mut usage = sample_user_usage_audit(
|
||||
"usage-dashboard-auth-fallback-1",
|
||||
"req-dashboard-auth-fallback-1",
|
||||
"user-auth-1",
|
||||
"gpt-5",
|
||||
"OpenAI",
|
||||
"completed",
|
||||
now - chrono::Duration::minutes(5),
|
||||
);
|
||||
usage.username = Some("stale-alice".to_string());
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
|
||||
|
||||
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
|
||||
start_auth_gateway_with_builder(|| {
|
||||
let wallet_repository =
|
||||
Arc::new(InMemoryWalletRepository::seed(vec![sample_auth_wallet(
|
||||
"user-auth-1",
|
||||
now,
|
||||
)]));
|
||||
let data_state = GatewayDataState::with_user_wallet_and_usage_for_tests(
|
||||
Arc::new(InMemoryUserReadRepository::seed_auth_users(Vec::<
|
||||
StoredUserAuthRecord,
|
||||
>::new(
|
||||
))),
|
||||
wallet_repository,
|
||||
usage_repository,
|
||||
);
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_auth_users_for_tests([user.clone()])
|
||||
.with_auth_sessions_for_tests([session])
|
||||
})
|
||||
.await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/dashboard/recent-requests?limit=5"
|
||||
))
|
||||
.header("authorization", format!("Bearer {access_token}"))
|
||||
.header(
|
||||
"x-client-device-id",
|
||||
"device-dashboard-recent-auth-fallback",
|
||||
)
|
||||
.header("user-agent", "AetherTest/1.0")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let requests = payload["requests"].as_array().expect("array");
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(requests[0]["id"], "usage-dashboard-auth-fallback-1");
|
||||
assert_eq!(requests[0]["user"], "alice");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_dashboard_provider_status_locally_without_proxying_upstream() {
|
||||
let now = Utc::now();
|
||||
|
||||
@@ -9,6 +9,16 @@ use super::{
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
fn large_request_body(stream: bool) -> String {
|
||||
let large_message = "x".repeat(128 * 1024);
|
||||
serde_json::to_string(&json!({
|
||||
"model": "gpt-5",
|
||||
"messages": [{"role": "user", "content": large_message}],
|
||||
"stream": stream
|
||||
}))
|
||||
.expect("request body should encode")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled() {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
@@ -271,6 +281,147 @@ async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arri
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_keeps_pending_sync_usage_lightweight_for_large_request_body() {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let execution_request_started = Arc::new(tokio::sync::Notify::new());
|
||||
let allow_execution_response = Arc::new(tokio::sync::Notify::new());
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(|_request: Request| async move { Json(json!({"ok": true})) }),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any({
|
||||
let execution_request_started = Arc::clone(&execution_request_started);
|
||||
let allow_execution_response = Arc::clone(&allow_execution_response);
|
||||
move |_request: Request| {
|
||||
let execution_request_started = Arc::clone(&execution_request_started);
|
||||
let allow_execution_response = Arc::clone(&allow_execution_response);
|
||||
async move {
|
||||
execution_request_started.notify_one();
|
||||
allow_execution_response.notified().await;
|
||||
Json(json!({
|
||||
"request_id": "req-usage-sync-large-pending-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "chatcmpl-usage-sync-large-pending-123",
|
||||
"usage": {
|
||||
"input_tokens": 3,
|
||||
"output_tokens": 5,
|
||||
"total_tokens": 8
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 45
|
||||
}
|
||||
}))
|
||||
}
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-usage-sync-large-pending")),
|
||||
sample_local_openai_auth_snapshot(
|
||||
"api-key-usage-sync-large-pending-123",
|
||||
"user-usage-sync-large-pending-123",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_local_openai_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_local_openai_provider()],
|
||||
vec![sample_local_openai_endpoint()],
|
||||
vec![sample_local_openai_key()],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
)
|
||||
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
});
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let request_task = tokio::spawn({
|
||||
let gateway_url = gateway_url.clone();
|
||||
async move {
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-usage-sync-large-pending",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "req-usage-sync-large-pending-123")
|
||||
.body(large_request_body(false))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
let status = response.status();
|
||||
let body = response.text().await.expect("body should read");
|
||||
(status, body)
|
||||
}
|
||||
});
|
||||
|
||||
execution_request_started.notified().await;
|
||||
|
||||
let mut pending = None;
|
||||
for _ in 0..50 {
|
||||
pending = usage_repository
|
||||
.find_by_request_id("req-usage-sync-large-pending-123")
|
||||
.await
|
||||
.expect("usage lookup should succeed");
|
||||
if pending
|
||||
.as_ref()
|
||||
.is_some_and(|stored| stored.status == "pending")
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
let pending = pending.expect("pending usage should be recorded before sync result resolves");
|
||||
assert_eq!(pending.status, "pending");
|
||||
assert!(pending.request_headers.is_none());
|
||||
assert!(pending.request_body.is_none());
|
||||
assert!(pending.provider_request_headers.is_none());
|
||||
assert!(pending.provider_request_body.is_none());
|
||||
assert!(pending.response_headers.is_none());
|
||||
assert!(pending.client_response_headers.is_none());
|
||||
|
||||
allow_execution_response.notify_one();
|
||||
|
||||
let (status, _body) = request_task.await.expect("request task should join");
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
async fn gateway_records_usage_for_execution_runtime_stream_when_runtime_enabled() {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
|
||||
@@ -540,3 +691,139 @@ async fn gateway_records_pending_usage_before_execution_runtime_stream_headers_a
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_keeps_pending_stream_usage_lightweight_for_large_request_body() {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let execution_request_started = Arc::new(tokio::sync::Notify::new());
|
||||
let allow_execution_response = Arc::new(tokio::sync::Notify::new());
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/gateway/report-stream",
|
||||
any(|_request: Request| async move { Json(json!({"ok": true})) }),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
any({
|
||||
let execution_request_started = Arc::clone(&execution_request_started);
|
||||
let allow_execution_response = Arc::clone(&allow_execution_response);
|
||||
move |_request: Request| {
|
||||
let execution_request_started = Arc::clone(&execution_request_started);
|
||||
let allow_execution_response = Arc::clone(&allow_execution_response);
|
||||
async move {
|
||||
execution_request_started.notify_one();
|
||||
allow_execution_response.notified().await;
|
||||
let frames = concat!(
|
||||
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-large-pending-123\\\",\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n",
|
||||
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":51,\"ttfb_ms\":19}}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(frames))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-stream-large-pending")),
|
||||
sample_local_openai_auth_snapshot(
|
||||
"api-key-usage-stream-large-pending-123",
|
||||
"user-usage-stream-large-pending-123",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_local_openai_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_local_openai_provider()],
|
||||
vec![sample_local_openai_endpoint()],
|
||||
vec![sample_local_openai_key()],
|
||||
));
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
usage_repository.clone(),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
)
|
||||
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
});
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let request_task = tokio::spawn({
|
||||
let gateway_url = gateway_url.clone();
|
||||
async move {
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-stream-large-pending",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "req-usage-stream-large-pending-123")
|
||||
.body(large_request_body(true))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
let status = response.status();
|
||||
let body = response.text().await.expect("stream body should read");
|
||||
(status, body)
|
||||
}
|
||||
});
|
||||
|
||||
execution_request_started.notified().await;
|
||||
|
||||
let mut pending = None;
|
||||
for _ in 0..50 {
|
||||
pending = usage_repository
|
||||
.find_by_request_id("req-usage-stream-large-pending-123")
|
||||
.await
|
||||
.expect("usage lookup should succeed");
|
||||
if pending
|
||||
.as_ref()
|
||||
.is_some_and(|stored| stored.status == "pending")
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
let pending = pending.expect("pending usage should be recorded before stream headers arrive");
|
||||
assert_eq!(pending.status, "pending");
|
||||
assert!(pending.request_headers.is_none());
|
||||
assert!(pending.request_body.is_none());
|
||||
assert!(pending.provider_request_headers.is_none());
|
||||
assert!(pending.provider_request_body.is_none());
|
||||
assert!(pending.response_headers.is_none());
|
||||
assert!(pending.client_response_headers.is_none());
|
||||
|
||||
allow_execution_response.notify_one();
|
||||
|
||||
let (status, _body) = request_task.await.expect("request task should join");
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
@@ -1,13 +1,28 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, hash_api_key, json,
|
||||
sample_local_openai_auth_snapshot, sample_local_openai_candidate_row,
|
||||
sample_local_openai_endpoint, sample_local_openai_key, sample_local_openai_provider,
|
||||
start_server, Arc, Body, GatewayDataState, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
|
||||
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
|
||||
InMemoryRequestCandidateRepository, InMemoryUsageReadRepository, Json, Mutex, Request,
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, Response, Router, StatusCode,
|
||||
UsageReadRepository, UsageRuntimeConfig, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override,
|
||||
encrypt_python_fernet_plaintext, hash_api_key, json, sample_local_openai_auth_snapshot,
|
||||
sample_local_openai_candidate_row, sample_local_openai_endpoint, sample_local_openai_key,
|
||||
sample_local_openai_provider, start_server, Arc, Body, GatewayDataState, HeaderValue,
|
||||
InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository,
|
||||
InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository,
|
||||
InMemoryUsageReadRepository, Json, Mutex, Request, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, Response, Router, StatusCode, StoredAuthApiKeySnapshot,
|
||||
StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider, StoredProviderModelMapping, UsageReadRepository,
|
||||
UsageRuntimeConfig, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
};
|
||||
use crate::constants::LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER;
|
||||
|
||||
fn deep_nested_metadata(levels: usize) -> serde_json::Value {
|
||||
let mut current = json!({"leaf": "value"});
|
||||
for depth in 0..levels {
|
||||
current = json!({
|
||||
"depth": depth,
|
||||
"child": current
|
||||
});
|
||||
}
|
||||
current
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when_usage_runtime_enabled(
|
||||
@@ -183,6 +198,157 @@ async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage() {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(|_request: Request| async move { Json(json!({"ok": true})) }),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"request_id": "trace-openai-chat-local-report-sync-deep-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "chatcmpl-local-report-sync-deep-123",
|
||||
"object": "chat.completion",
|
||||
"model": "gpt-5-upstream",
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 2,
|
||||
"completion_tokens": 3,
|
||||
"total_tokens": 5
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 25
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-local-report-sync-deep")),
|
||||
sample_local_openai_auth_snapshot(
|
||||
"api-key-openai-usage-local-deep-1",
|
||||
"user-openai-usage-local-deep-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_local_openai_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_local_openai_provider()],
|
||||
vec![sample_local_openai_endpoint()],
|
||||
vec![sample_local_openai_key()],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
)
|
||||
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
});
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let request_body = serde_json::to_string(&json!({
|
||||
"model": "gpt-5",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": "x".repeat(128 * 1024)
|
||||
}],
|
||||
"metadata": deep_nested_metadata(96)
|
||||
}))
|
||||
.expect("request should encode");
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-local-report-sync-deep",
|
||||
)
|
||||
.header(
|
||||
TRACE_ID_HEADER,
|
||||
"trace-openai-chat-local-report-sync-deep-123",
|
||||
)
|
||||
.body(request_body)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let mut stored_usage = None;
|
||||
for _ in 0..50 {
|
||||
stored_usage = usage_repository
|
||||
.find_by_request_id("trace-openai-chat-local-report-sync-deep-123")
|
||||
.await
|
||||
.expect("usage lookup should succeed");
|
||||
if stored_usage.is_some() {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
let stored_usage = stored_usage.expect("usage should be recorded");
|
||||
assert_eq!(stored_usage.status, "completed");
|
||||
assert_eq!(stored_usage.total_tokens, 5);
|
||||
assert_eq!(
|
||||
stored_usage
|
||||
.request_body
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("messages"))
|
||||
.and_then(|value| value.as_array())
|
||||
.and_then(|messages| messages.first())
|
||||
.and_then(|value| value.get("content"))
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::len),
|
||||
Some(128 * 1024)
|
||||
);
|
||||
assert_eq!(
|
||||
stored_usage
|
||||
.request_body
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("child"))
|
||||
.and_then(|value| value.get("child"))
|
||||
.and_then(|value| value.get("child"))
|
||||
.and_then(|value| value.get("child"))
|
||||
.and_then(|value| value.get("child"))
|
||||
.and_then(|value| value.as_object())
|
||||
.map(|value| value.contains_key("depth")),
|
||||
Some(true)
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exhaust_after_retryable_sync_failure(
|
||||
) {
|
||||
@@ -389,6 +555,220 @@ async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exha
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_records_failed_usage_for_claude_runtime_miss_without_execution_exhaustion() {
|
||||
fn sample_claude_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:chat"])),
|
||||
Some(serde_json::json!(["claude-sonnet-4-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:chat"])),
|
||||
Some(serde_json::json!(["claude-sonnet-4-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let report_hits = Arc::new(Mutex::new(0usize));
|
||||
let report_hits_clone = Arc::clone(&report_hits);
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/decision-sync",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |_request: Request| {
|
||||
let report_hits_inner = Arc::clone(&report_hits_clone);
|
||||
async move {
|
||||
*report_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/messages",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new();
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-claude-runtime-miss-usage")),
|
||||
sample_claude_auth_snapshot(
|
||||
"api-key-claude-runtime-miss-usage-1",
|
||||
"user-claude-runtime-miss-usage-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository = Arc::new(
|
||||
InMemoryMinimalCandidateSelectionReadRepository::seed(vec![]),
|
||||
);
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
)
|
||||
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
});
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages?beta=true"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("x-api-key", "sk-client-claude-runtime-miss-usage")
|
||||
.header("anthropic-version", "2023-06-01")
|
||||
.header(TRACE_ID_HEADER, "trace-claude-runtime-miss-usage-123")
|
||||
.body("{\"model\":\"claude-sonnet-4-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should complete");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("candidate_list_empty")
|
||||
);
|
||||
let body_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(body_json["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
body_json["error"]["message"],
|
||||
"没有可用的提供商支持模型 claude-sonnet-4-5 的同步请求"
|
||||
);
|
||||
|
||||
let mut stored_usage = None;
|
||||
for _ in 0..50 {
|
||||
stored_usage = usage_repository
|
||||
.find_by_request_id("trace-claude-runtime-miss-usage-123")
|
||||
.await
|
||||
.expect("usage lookup should succeed");
|
||||
if stored_usage.is_some() {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
let stored_usage = stored_usage.expect("failed usage should be recorded");
|
||||
assert_eq!(stored_usage.status, "failed");
|
||||
assert_eq!(stored_usage.billing_status, "void");
|
||||
assert_eq!(stored_usage.status_code, Some(503));
|
||||
assert_eq!(stored_usage.error_category.as_deref(), Some("server_error"));
|
||||
assert_eq!(
|
||||
stored_usage.user_id.as_deref(),
|
||||
Some("user-claude-runtime-miss-usage-1")
|
||||
);
|
||||
assert_eq!(stored_usage.provider_name, "claude");
|
||||
assert_eq!(stored_usage.model, "claude-sonnet-4-5");
|
||||
assert_eq!(stored_usage.api_format.as_deref(), Some("claude:chat"));
|
||||
assert_eq!(
|
||||
stored_usage.routing_execution_path(),
|
||||
Some("local_execution_runtime_miss")
|
||||
);
|
||||
assert_eq!(
|
||||
stored_usage.routing_local_execution_runtime_miss_reason(),
|
||||
Some("candidate_list_empty")
|
||||
);
|
||||
assert_eq!(stored_usage.routing_route_family(), Some("claude"));
|
||||
assert_eq!(stored_usage.routing_route_kind(), Some("chat"));
|
||||
assert_eq!(
|
||||
stored_usage
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("trace_id"))
|
||||
.and_then(|value| value.as_str()),
|
||||
Some("trace-claude-runtime-miss-usage-123")
|
||||
);
|
||||
assert_eq!(
|
||||
stored_usage
|
||||
.client_response_body
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("error"))
|
||||
.and_then(|value| value.get("type"))
|
||||
.and_then(|value| value.as_str()),
|
||||
Some("http_error")
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-claude-runtime-miss-usage-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert!(stored_candidates.is_empty());
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert_eq!(*report_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_when_usage_runtime_enabled(
|
||||
) {
|
||||
@@ -559,3 +939,552 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_skipped() {
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["gpt-5.4"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
None,
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["gpt-5.4"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-claude-cli-usage-local-miss-1".to_string(),
|
||||
provider_name: "RightCode".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-claude-cli-usage-local-miss-1".to_string(),
|
||||
endpoint_api_format: "openai:cli".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-claude-cli-usage-local-miss-1".to_string(),
|
||||
key_name: "codex".to_string(),
|
||||
key_auth_type: "bearer".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:cli".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:cli": 1})),
|
||||
model_id: "model-claude-cli-usage-local-miss-1".to_string(),
|
||||
global_model_id: "global-model-claude-cli-usage-local-miss-1".to_string(),
|
||||
global_model_name: "gpt-5.4".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-5.4".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5.4".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:cli".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-claude-cli-usage-local-miss-1".to_string(),
|
||||
"RightCode".to_string(),
|
||||
Some("https://right.codes".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-claude-cli-usage-local-miss-1".to_string(),
|
||||
"provider-claude-cli-usage-local-miss-1".to_string(),
|
||||
"openai:cli".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://right.codes/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some("/v1/messages".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-claude-cli-usage-local-miss-1".to_string(),
|
||||
"provider-claude-cli-usage-local-miss-1".to_string(),
|
||||
"codex".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["openai:cli"])),
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"sk-upstream-openai-cli-usage-local-miss",
|
||||
)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"openai:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/v1/messages",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
let execution_runtime = Router::new();
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-claude-cli-usage-local-miss")),
|
||||
sample_auth_snapshot(
|
||||
"api-key-claude-cli-usage-local-miss-1",
|
||||
"user-claude-cli-usage-local-miss-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
)
|
||||
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
});
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages?beta=true"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-claude-cli-usage-local-miss",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-claude-cli-usage-local-miss-123")
|
||||
.body("{\"model\":\"gpt-5.4\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should complete");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("all_candidates_skipped")
|
||||
);
|
||||
let body_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(body_json["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
body_json["error"]["message"],
|
||||
"没有可用的提供商支持模型 gpt-5.4 的同步请求"
|
||||
);
|
||||
|
||||
let mut stored_usage = None;
|
||||
for _ in 0..50 {
|
||||
stored_usage = usage_repository
|
||||
.find_by_request_id("trace-claude-cli-usage-local-miss-123")
|
||||
.await
|
||||
.expect("usage lookup should succeed");
|
||||
if stored_usage.is_some() {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
let stored_usage = stored_usage.expect("failed usage should be recorded");
|
||||
assert_eq!(stored_usage.status, "failed");
|
||||
assert_eq!(stored_usage.billing_status, "void");
|
||||
assert_eq!(stored_usage.status_code, Some(503));
|
||||
assert_eq!(stored_usage.error_category.as_deref(), Some("server_error"));
|
||||
assert_eq!(
|
||||
stored_usage.user_id.as_deref(),
|
||||
Some("user-claude-cli-usage-local-miss-1")
|
||||
);
|
||||
assert_eq!(stored_usage.provider_name, "RightCode");
|
||||
assert_eq!(stored_usage.model, "gpt-5.4");
|
||||
assert_eq!(stored_usage.api_format.as_deref(), Some("claude:cli"));
|
||||
assert_eq!(
|
||||
stored_usage.endpoint_api_format.as_deref(),
|
||||
Some("openai:cli")
|
||||
);
|
||||
assert_eq!(stored_usage.routing_key_name(), Some("codex"));
|
||||
assert_eq!(stored_usage.routing_planner_kind(), Some("claude_cli_sync"));
|
||||
assert_eq!(stored_usage.routing_route_family(), Some("claude"));
|
||||
assert_eq!(stored_usage.routing_route_kind(), Some("cli"));
|
||||
assert_eq!(
|
||||
stored_usage.routing_execution_path(),
|
||||
Some("local_execution_runtime_miss")
|
||||
);
|
||||
assert_eq!(
|
||||
stored_usage.routing_local_execution_runtime_miss_reason(),
|
||||
Some("all_candidates_skipped")
|
||||
);
|
||||
assert_eq!(
|
||||
stored_usage.error_message.as_deref(),
|
||||
Some("没有可用的提供商支持模型 gpt-5.4 的同步请求")
|
||||
);
|
||||
assert_eq!(
|
||||
stored_usage
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("trace_id"))
|
||||
.and_then(|value| value.as_str()),
|
||||
Some("trace-claude-cli-usage-local-miss-123")
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-claude-cli-usage-local-miss-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Skipped);
|
||||
assert_eq!(
|
||||
stored_candidates[0].skip_reason.as_deref(),
|
||||
Some("format_conversion_disabled")
|
||||
);
|
||||
assert_eq!(
|
||||
stored_usage.routing_candidate_id(),
|
||||
Some(stored_candidates[0].id.as_str())
|
||||
);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_keeps_failed_usage_request_capture_lightweight_for_large_local_claude_cli_runtime_miss(
|
||||
) {
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["gpt-5.4"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
None,
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["gpt-5.4"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-claude-cli-usage-local-miss-large-1".to_string(),
|
||||
provider_name: "RightCode".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-claude-cli-usage-local-miss-large-1".to_string(),
|
||||
endpoint_api_format: "openai:cli".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-claude-cli-usage-local-miss-large-1".to_string(),
|
||||
key_name: "codex".to_string(),
|
||||
key_auth_type: "bearer".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:cli".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:cli": 1})),
|
||||
model_id: "model-claude-cli-usage-local-miss-large-1".to_string(),
|
||||
global_model_id: "global-model-claude-cli-usage-local-miss-large-1".to_string(),
|
||||
global_model_name: "gpt-5.4".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-5.4".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5.4".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:cli".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-claude-cli-usage-local-miss-large-1".to_string(),
|
||||
"RightCode".to_string(),
|
||||
Some("https://right.codes".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-claude-cli-usage-local-miss-large-1".to_string(),
|
||||
"provider-claude-cli-usage-local-miss-large-1".to_string(),
|
||||
"openai:cli".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://right.codes/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some("/v1/messages".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-claude-cli-usage-local-miss-large-1".to_string(),
|
||||
"provider-claude-cli-usage-local-miss-large-1".to_string(),
|
||||
"codex".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["openai:cli"])),
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"sk-upstream-openai-cli-usage-local-miss-large",
|
||||
)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"openai:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
|
||||
let execution_runtime = Router::new();
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-claude-cli-usage-local-miss-large")),
|
||||
sample_auth_snapshot(
|
||||
"api-key-claude-cli-usage-local-miss-large-1",
|
||||
"user-claude-cli-usage-local-miss-large-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
)
|
||||
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
..UsageRuntimeConfig::default()
|
||||
});
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let request_body = serde_json::to_string(&json!({
|
||||
"model": "gpt-5.4",
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": "x".repeat(128 * 1024)
|
||||
}],
|
||||
"metadata": deep_nested_metadata(96)
|
||||
}))
|
||||
.expect("request should encode");
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages?beta=true"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-claude-cli-usage-local-miss-large",
|
||||
)
|
||||
.header(
|
||||
TRACE_ID_HEADER,
|
||||
"trace-claude-cli-usage-local-miss-large-123",
|
||||
)
|
||||
.body(request_body)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should complete");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("all_candidates_skipped")
|
||||
);
|
||||
|
||||
let mut stored_usage = None;
|
||||
for _ in 0..50 {
|
||||
stored_usage = usage_repository
|
||||
.find_by_request_id("trace-claude-cli-usage-local-miss-large-123")
|
||||
.await
|
||||
.expect("usage lookup should succeed");
|
||||
if stored_usage.is_some() {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
let stored_usage = stored_usage.expect("failed usage should be recorded");
|
||||
assert_eq!(stored_usage.status, "failed");
|
||||
assert!(stored_usage.request_body.is_none());
|
||||
assert_eq!(
|
||||
stored_usage
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("trace_id"))
|
||||
.and_then(|value| value.as_str()),
|
||||
Some("trace-claude-cli-usage-local-miss-large-123")
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-claude-cli-usage-local-miss-large-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Skipped);
|
||||
assert_eq!(
|
||||
stored_candidates[0].skip_reason.as_deref(),
|
||||
Some("format_conversion_disabled")
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user