mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 18:37:46 +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:
@@ -8,6 +8,7 @@ pub(crate) mod transport;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::{Response, Uri};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
|
||||
|
||||
@@ -58,6 +59,10 @@ pub(crate) fn collect_control_headers(
|
||||
crate::headers::collect_control_headers(headers)
|
||||
}
|
||||
|
||||
pub(crate) fn build_report_context_original_request_echo(body_json: &Value) -> Option<Value> {
|
||||
(!body_json.is_null()).then(|| body_json.clone())
|
||||
}
|
||||
|
||||
pub(crate) fn is_json_request(headers: &http::HeaderMap) -> bool {
|
||||
crate::headers::is_json_request(headers)
|
||||
}
|
||||
@@ -99,3 +104,26 @@ pub(crate) fn maybe_build_local_sync_finalize_response(
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
crate::execution_runtime::maybe_build_local_sync_finalize_response(trace_id, decision, payload)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::build_report_context_original_request_echo;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn build_report_context_original_request_echo_preserves_full_request_body() {
|
||||
let body = json!({
|
||||
"messages": [{"role": "user", "content": "large payload should be omitted"}],
|
||||
"service_tier": "default",
|
||||
"instructions": "Be concise.",
|
||||
"thinking": {"type": "enabled", "budget_tokens": 512},
|
||||
"metadata": {"trace": "keep"},
|
||||
"body_bytes_b64": "aGVsbG8=",
|
||||
});
|
||||
|
||||
let echo =
|
||||
build_report_context_original_request_echo(&body).expect("echo should be produced");
|
||||
|
||||
assert_eq!(echo, body);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use tracing::debug;
|
||||
use url::Url;
|
||||
|
||||
use crate::ai_pipeline::planner::common::{
|
||||
@@ -155,6 +156,25 @@ async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
|
||||
} else {
|
||||
ConversionMode::Bidirectional
|
||||
};
|
||||
let upstream_base_url = infer_upstream_base_url(&follow_up.plan.url);
|
||||
|
||||
debug!(
|
||||
event_name = "local_video_follow_up_sync_decision_payload_built",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id,
|
||||
request_id = %trace_id,
|
||||
candidate_id = ?follow_up.plan.candidate_id,
|
||||
provider_id = %follow_up.plan.provider_id,
|
||||
endpoint_id = %follow_up.plan.endpoint_id,
|
||||
key_id = %follow_up.plan.key_id,
|
||||
plan_kind,
|
||||
downstream_path = %parts.uri.path(),
|
||||
provider_api_format = %follow_up.plan.provider_api_format,
|
||||
client_api_format = %follow_up.plan.client_api_format,
|
||||
upstream_base_url = ?upstream_base_url,
|
||||
upstream_url = %follow_up.plan.url,
|
||||
"gateway built local video follow-up sync decision payload"
|
||||
);
|
||||
|
||||
Ok(Some(GatewayControlSyncDecisionResponse {
|
||||
action: EXECUTION_RUNTIME_SYNC_DECISION_ACTION.to_string(),
|
||||
@@ -167,7 +187,7 @@ async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
|
||||
provider_id: Some(follow_up.plan.provider_id.clone()),
|
||||
endpoint_id: Some(follow_up.plan.endpoint_id.clone()),
|
||||
key_id: Some(follow_up.plan.key_id.clone()),
|
||||
upstream_base_url: infer_upstream_base_url(&follow_up.plan.url),
|
||||
upstream_base_url,
|
||||
upstream_url: Some(follow_up.plan.url.clone()),
|
||||
provider_request_method: Some(follow_up.plan.method.clone()),
|
||||
auth_header: auth_pair.as_ref().map(|(name, _)| name.clone()),
|
||||
@@ -219,5 +239,79 @@ fn infer_upstream_base_url(upstream_url: &str) -> Option<String> {
|
||||
base.push(':');
|
||||
base.push_str(port.to_string().as_str());
|
||||
}
|
||||
let base_path = infer_upstream_base_path(parsed.path());
|
||||
if !base_path.is_empty() {
|
||||
base.push_str(base_path);
|
||||
}
|
||||
Some(base)
|
||||
}
|
||||
|
||||
fn infer_upstream_base_path(path: &str) -> &str {
|
||||
let trimmed = path.trim_end_matches('/');
|
||||
if trimmed.is_empty() || trimmed == "/" {
|
||||
return "";
|
||||
}
|
||||
|
||||
for suffix in [
|
||||
"/responses/compact",
|
||||
"/responses",
|
||||
"/chat/completions",
|
||||
"/messages",
|
||||
] {
|
||||
if let Some(prefix) = trimmed.strip_suffix(suffix) {
|
||||
return normalize_inferred_base_path(prefix);
|
||||
}
|
||||
}
|
||||
|
||||
for marker in ["/v1/videos", "/v1beta/"] {
|
||||
if let Some((prefix, _)) = trimmed.split_once(marker) {
|
||||
return normalize_inferred_base_path(prefix);
|
||||
}
|
||||
}
|
||||
|
||||
normalize_inferred_base_path(trimmed)
|
||||
}
|
||||
|
||||
fn normalize_inferred_base_path(path: &str) -> &str {
|
||||
let trimmed = path.trim_end_matches('/');
|
||||
if trimmed.is_empty() || trimmed == "/" {
|
||||
""
|
||||
} else {
|
||||
trimmed
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::infer_upstream_base_url;
|
||||
|
||||
#[test]
|
||||
fn infer_upstream_base_url_preserves_codex_base_path() {
|
||||
assert_eq!(
|
||||
infer_upstream_base_url("https://tiger.bookapi.cc/codex/responses").as_deref(),
|
||||
Some("https://tiger.bookapi.cc/codex")
|
||||
);
|
||||
assert_eq!(
|
||||
infer_upstream_base_url("https://chatgpt.com/backend-api/codex/responses").as_deref(),
|
||||
Some("https://chatgpt.com/backend-api/codex")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_upstream_base_url_preserves_nested_v1_prefix() {
|
||||
assert_eq!(
|
||||
infer_upstream_base_url("https://api.openai.example/custom/v1/chat/completions?mode=1")
|
||||
.as_deref(),
|
||||
Some("https://api.openai.example/custom/v1")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_upstream_base_url_strips_video_operation_path() {
|
||||
assert_eq!(
|
||||
infer_upstream_base_url("https://video.example/nested/v1/videos/task-123/content")
|
||||
.as_deref(),
|
||||
Some("https://video.example/nested")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
use crate::{
|
||||
AppState, GatewayControlSyncDecisionResponse, GatewayError, LocalExecutionRuntimeMissDiagnostic,
|
||||
};
|
||||
|
||||
use super::super::plans::{resolve_stream_spec, resolve_sync_spec};
|
||||
use super::candidates::{
|
||||
@@ -8,6 +10,52 @@ use super::candidates::{
|
||||
};
|
||||
use super::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
|
||||
|
||||
fn extract_requested_model(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
spec: crate::ai_pipeline::LocalSameFormatProviderSpec,
|
||||
) -> Option<String> {
|
||||
match spec.family {
|
||||
crate::ai_pipeline::LocalSameFormatProviderFamily::Standard => body_json
|
||||
.get("model")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
crate::ai_pipeline::LocalSameFormatProviderFamily::Gemini => {
|
||||
let marker = "/models/";
|
||||
let start = parts.uri.path().find(marker)? + marker.len();
|
||||
let tail = &parts.uri.path()[start..];
|
||||
let end = tail.find(':').unwrap_or(tail.len());
|
||||
let model = tail[..end].trim();
|
||||
if model.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(model.to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_local_same_format_miss_diagnostic(
|
||||
decision: &GatewayControlDecision,
|
||||
spec: crate::ai_pipeline::LocalSameFormatProviderSpec,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) -> LocalExecutionRuntimeMissDiagnostic {
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
reason: reason.to_string(),
|
||||
route_family: decision.route_family.clone(),
|
||||
route_kind: decision.route_kind.clone(),
|
||||
public_path: Some(decision.public_path.clone()),
|
||||
plan_kind: Some(spec.decision_kind.to_string()),
|
||||
requested_model: requested_model.map(ToOwned::to_owned),
|
||||
candidate_count: None,
|
||||
skipped_candidate_count: None,
|
||||
skip_reasons: std::collections::BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
@@ -25,12 +73,42 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
|
||||
)
|
||||
.await
|
||||
else {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_same_format_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
extract_requested_model(parts, body_json, spec).as_deref(),
|
||||
"decision_input_unavailable",
|
||||
),
|
||||
);
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let attempts =
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_same_format_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
),
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
|
||||
.await?;
|
||||
let preserve_existing_candidate_signal = candidate_count == 0
|
||||
&& state.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id);
|
||||
if !preserve_existing_candidate_signal {
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else {
|
||||
"candidate_evaluation_incomplete".to_string()
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) =
|
||||
@@ -43,6 +121,18 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
|
||||
}
|
||||
}
|
||||
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
|
||||
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else if skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_sync_plans".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
@@ -63,12 +153,42 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
|
||||
)
|
||||
.await
|
||||
else {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_same_format_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
extract_requested_model(parts, body_json, spec).as_deref(),
|
||||
"decision_input_unavailable",
|
||||
),
|
||||
);
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let attempts =
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_same_format_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
),
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
|
||||
.await?;
|
||||
let preserve_existing_candidate_signal = candidate_count == 0
|
||||
&& state.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id);
|
||||
if !preserve_existing_candidate_signal {
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else {
|
||||
"candidate_evaluation_incomplete".to_string()
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) =
|
||||
@@ -81,5 +201,17 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
|
||||
}
|
||||
}
|
||||
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
|
||||
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else if skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_stream_plans".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
+3
-2
@@ -85,7 +85,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
trace_id: &str,
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Result<Vec<LocalSameFormatProviderCandidateAttempt>, GatewayError> {
|
||||
) -> Result<(Vec<LocalSameFormatProviderCandidateAttempt>, usize), GatewayError> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let candidates = planner_state
|
||||
.list_selectable_candidates(
|
||||
@@ -104,6 +104,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
input.required_capabilities.as_ref(),
|
||||
)
|
||||
.await;
|
||||
let candidate_count = candidates.len();
|
||||
|
||||
let created_at_unix_ms = current_unix_ms();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
@@ -160,5 +161,5 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
});
|
||||
}
|
||||
|
||||
Ok(attempts)
|
||||
Ok((attempts, candidate_count))
|
||||
}
|
||||
|
||||
@@ -287,9 +287,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
"upstream_url": upstream_url,
|
||||
"provider_request_method": serde_json::Value::Null,
|
||||
"provider_request_headers": provider_request_headers,
|
||||
"provider_request_body": provider_request_body,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
|
||||
"has_envelope": is_kiro || is_antigravity,
|
||||
"envelope_name": if is_kiro {
|
||||
Some(KIRO_ENVELOPE_NAME)
|
||||
@@ -359,6 +358,13 @@ pub(super) async fn mark_skipped_local_same_format_provider_candidate(
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
*diagnostic
|
||||
.skip_reasons
|
||||
.entry(skip_reason.to_string())
|
||||
.or_insert(0) += 1;
|
||||
*diagnostic.skipped_candidate_count.get_or_insert(0) += 1;
|
||||
});
|
||||
PlannerAppState::new(state)
|
||||
.persist_skipped_local_candidate(
|
||||
trace_id,
|
||||
|
||||
@@ -16,6 +16,53 @@ use crate::ai_pipeline::planner::plan_builders::{
|
||||
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
|
||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||
};
|
||||
use crate::LocalExecutionRuntimeMissDiagnostic;
|
||||
|
||||
fn extract_requested_model(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<String> {
|
||||
match spec.family {
|
||||
LocalSameFormatProviderFamily::Standard => body_json
|
||||
.get("model")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
LocalSameFormatProviderFamily::Gemini => {
|
||||
let marker = "/models/";
|
||||
let start = parts.uri.path().find(marker)? + marker.len();
|
||||
let tail = &parts.uri.path()[start..];
|
||||
let end = tail.find(':').unwrap_or(tail.len());
|
||||
let model = tail[..end].trim();
|
||||
if model.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(model.to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_local_same_format_miss_diagnostic(
|
||||
decision: &GatewayControlDecision,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) -> LocalExecutionRuntimeMissDiagnostic {
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
reason: reason.to_string(),
|
||||
route_family: decision.route_family.clone(),
|
||||
route_kind: decision.route_kind.clone(),
|
||||
public_path: Some(decision.public_path.clone()),
|
||||
plan_kind: Some(spec.decision_kind.to_string()),
|
||||
requested_model: requested_model.map(ToOwned::to_owned),
|
||||
candidate_count: None,
|
||||
skipped_candidate_count: None,
|
||||
skip_reasons: std::collections::BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
state: &AppState,
|
||||
@@ -30,12 +77,44 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
)
|
||||
.await
|
||||
else {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_same_format_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
extract_requested_model(parts, body_json, spec).as_deref(),
|
||||
"decision_input_unavailable",
|
||||
),
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let attempts =
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_same_format_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
),
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
|
||||
.await?;
|
||||
let preserve_existing_candidate_signal = candidate_count == 0
|
||||
&& state.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id);
|
||||
if !preserve_existing_candidate_signal {
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else {
|
||||
"candidate_evaluation_incomplete".to_string()
|
||||
};
|
||||
});
|
||||
}
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
@@ -70,6 +149,16 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
}
|
||||
}
|
||||
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
|
||||
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
|
||||
diagnostic.reason = if candidate_count > 0 && skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_sync_plans".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
@@ -86,12 +175,44 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
)
|
||||
.await
|
||||
else {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_same_format_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
extract_requested_model(parts, body_json, spec).as_deref(),
|
||||
"decision_input_unavailable",
|
||||
),
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let attempts =
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_same_format_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
),
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
|
||||
.await?;
|
||||
let preserve_existing_candidate_signal = candidate_count == 0
|
||||
&& state.local_execution_runtime_miss_diagnostic_has_candidate_signal(trace_id);
|
||||
if !preserve_existing_candidate_signal {
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else {
|
||||
"candidate_evaluation_incomplete".to_string()
|
||||
};
|
||||
});
|
||||
}
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
@@ -126,5 +247,15 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
}
|
||||
}
|
||||
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
|
||||
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
|
||||
diagnostic.reason = if candidate_count > 0 && skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_stream_plans".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
@@ -297,7 +297,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
"provider_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"client_api_format": GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": original_request_body,
|
||||
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(&original_request_body),
|
||||
"has_envelope": false,
|
||||
"needs_conversion": false,
|
||||
})),
|
||||
|
||||
@@ -263,7 +263,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
||||
"client_api_format": spec.api_format,
|
||||
"mapped_model": mapped_model,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
|
||||
"has_envelope": false,
|
||||
"needs_conversion": false,
|
||||
})),
|
||||
|
||||
@@ -6,7 +6,9 @@ use crate::ai_pipeline::planner::plan_builders::{
|
||||
LocalStreamPlanAndReport, LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::GatewayControlDecision;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
use crate::{
|
||||
AppState, GatewayControlSyncDecisionResponse, GatewayError, LocalExecutionRuntimeMissDiagnostic,
|
||||
};
|
||||
|
||||
use super::candidates::{
|
||||
materialize_local_standard_candidate_attempts, resolve_local_standard_decision_input,
|
||||
@@ -14,6 +16,52 @@ use super::candidates::{
|
||||
use super::payload::maybe_build_local_standard_decision_payload_for_candidate;
|
||||
use super::{LocalStandardSourceFamily, LocalStandardSpec};
|
||||
|
||||
fn extract_requested_model(
|
||||
parts: &http::request::Parts,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Option<String> {
|
||||
match spec.family {
|
||||
LocalStandardSourceFamily::Standard => body_json
|
||||
.get("model")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
LocalStandardSourceFamily::Gemini => {
|
||||
let marker = "/models/";
|
||||
let start = parts.uri.path().find(marker)? + marker.len();
|
||||
let tail = &parts.uri.path()[start..];
|
||||
let end = tail.find(':').unwrap_or(tail.len());
|
||||
let model = tail[..end].trim();
|
||||
if model.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(model.to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_local_standard_miss_diagnostic(
|
||||
decision: &GatewayControlDecision,
|
||||
spec: LocalStandardSpec,
|
||||
requested_model: Option<&str>,
|
||||
reason: &str,
|
||||
) -> LocalExecutionRuntimeMissDiagnostic {
|
||||
LocalExecutionRuntimeMissDiagnostic {
|
||||
reason: reason.to_string(),
|
||||
route_family: decision.route_family.clone(),
|
||||
route_kind: decision.route_kind.clone(),
|
||||
public_path: Some(decision.public_path.clone()),
|
||||
plan_kind: Some(spec.decision_kind.to_string()),
|
||||
requested_model: requested_model.map(ToOwned::to_owned),
|
||||
candidate_count: None,
|
||||
skipped_candidate_count: None,
|
||||
skip_reasons: std::collections::BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_sync_via_standard_family_payload(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
@@ -34,8 +82,25 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let attempts =
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_standard_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
),
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else {
|
||||
"candidate_evaluation_incomplete".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
@@ -47,6 +112,18 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
|
||||
}
|
||||
}
|
||||
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
|
||||
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else if skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_sync_plans".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
@@ -70,8 +147,25 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let attempts =
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_standard_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
),
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else {
|
||||
"candidate_evaluation_incomplete".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
for attempt in attempts {
|
||||
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
@@ -83,6 +177,18 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
|
||||
}
|
||||
}
|
||||
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
|
||||
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else if skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_stream_plans".to_string()
|
||||
};
|
||||
});
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
@@ -98,10 +204,39 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
else {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_standard_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
extract_requested_model(parts, body_json, spec).as_deref(),
|
||||
"decision_input_unavailable",
|
||||
),
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let attempts =
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_standard_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
),
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else {
|
||||
"candidate_evaluation_incomplete".to_string()
|
||||
};
|
||||
});
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
@@ -132,6 +267,15 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
}
|
||||
}
|
||||
}
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
|
||||
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
|
||||
diagnostic.reason = if candidate_count > 0 && skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_sync_plans".to_string()
|
||||
};
|
||||
});
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
@@ -147,10 +291,39 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
else {
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_standard_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
extract_requested_model(parts, body_json, spec).as_deref(),
|
||||
"decision_input_unavailable",
|
||||
),
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let attempts =
|
||||
state.set_local_execution_runtime_miss_diagnostic(
|
||||
trace_id,
|
||||
build_local_standard_miss_diagnostic(
|
||||
decision,
|
||||
spec,
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
),
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
diagnostic.candidate_count = Some(candidate_count);
|
||||
diagnostic.reason = if candidate_count == 0 {
|
||||
"candidate_list_empty".to_string()
|
||||
} else {
|
||||
"candidate_evaluation_incomplete".to_string()
|
||||
};
|
||||
});
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
@@ -181,5 +354,14 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
}
|
||||
}
|
||||
}
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
let candidate_count = diagnostic.candidate_count.unwrap_or(0);
|
||||
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
|
||||
diagnostic.reason = if candidate_count > 0 && skipped_candidate_count >= candidate_count {
|
||||
"all_candidates_skipped".to_string()
|
||||
} else {
|
||||
"no_local_stream_plans".to_string()
|
||||
};
|
||||
});
|
||||
Ok(plans)
|
||||
}
|
||||
|
||||
@@ -91,7 +91,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
trace_id: &str,
|
||||
input: &LocalStandardDecisionInput,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Result<Vec<LocalStandardCandidateAttempt>, GatewayError> {
|
||||
) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let mut seen_candidates = BTreeSet::new();
|
||||
let mut candidates = Vec::new();
|
||||
@@ -144,6 +144,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
input.required_capabilities.as_ref(),
|
||||
)
|
||||
.await;
|
||||
let candidate_count = candidates.len();
|
||||
|
||||
let created_at_unix_ms = current_unix_ms();
|
||||
let mut attempts = Vec::with_capacity(candidates.len());
|
||||
@@ -257,7 +258,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
});
|
||||
}
|
||||
|
||||
Ok(attempts)
|
||||
Ok((attempts, candidate_count))
|
||||
}
|
||||
|
||||
fn auth_snapshot_allows_cross_format_candidate(
|
||||
|
||||
@@ -358,9 +358,8 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
"upstream_url": upstream_url,
|
||||
"provider_request_method": serde_json::Value::Null,
|
||||
"provider_request_headers": provider_request_headers,
|
||||
"provider_request_body": provider_request_body,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
|
||||
"has_envelope": false,
|
||||
"needs_conversion": true,
|
||||
}),
|
||||
@@ -384,6 +383,13 @@ pub(super) async fn mark_skipped_local_standard_candidate(
|
||||
candidate_id: &str,
|
||||
skip_reason: &'static str,
|
||||
) {
|
||||
state.mutate_local_execution_runtime_miss_diagnostic(trace_id, |diagnostic| {
|
||||
*diagnostic
|
||||
.skip_reasons
|
||||
.entry(skip_reason.to_string())
|
||||
.or_insert(0) += 1;
|
||||
*diagnostic.skipped_candidate_count.get_or_insert(0) += 1;
|
||||
});
|
||||
PlannerAppState::new(state)
|
||||
.persist_skipped_local_candidate(
|
||||
trace_id,
|
||||
|
||||
@@ -1,6 +1,14 @@
|
||||
use aether_provider_transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
};
|
||||
use http::Request;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{build_cross_format_openai_cli_request_body, build_local_openai_cli_request_body};
|
||||
use super::{
|
||||
build_cross_format_openai_cli_request_body, build_local_openai_cli_request_body,
|
||||
build_local_openai_cli_upstream_url,
|
||||
};
|
||||
|
||||
fn object_keys(value: &Value) -> Vec<&str> {
|
||||
value
|
||||
@@ -11,6 +19,59 @@ fn object_keys(value: &Value) -> Vec<&str> {
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sample_transport(base_url: &str, api_format: &str) -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-codex".to_string(),
|
||||
name: "codex".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: false,
|
||||
concurrent_limit: None,
|
||||
max_retries: None,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: None,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-codex".to_string(),
|
||||
provider_id: "provider-codex".to_string(),
|
||||
api_format: api_format.to_string(),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
is_active: true,
|
||||
base_url: base_url.to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: None,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-codex".to_string(),
|
||||
provider_id: "provider-codex".to_string(),
|
||||
name: "oauth".to_string(),
|
||||
auth_type: "oauth".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec![api_format.to_string()]),
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
decrypted_api_key: "__placeholder__".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_openai_chat_cross_format_request_body_from_openai_cli_source() {
|
||||
let body_json = json!({
|
||||
@@ -103,6 +164,25 @@ fn local_openai_compact_wrapper_strips_store_for_same_format_requests() {
|
||||
assert!(provider_request_body.get("store").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_openai_cli_upstream_url_preserves_codex_base_path() {
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.body(())
|
||||
.expect("request should build");
|
||||
let (parts, _) = request.into_parts();
|
||||
|
||||
let upstream_url = build_local_openai_cli_upstream_url(
|
||||
&parts,
|
||||
&sample_transport("https://tiger.bookapi.cc/codex", "openai:cli"),
|
||||
false,
|
||||
)
|
||||
.expect("openai cli upstream url should build");
|
||||
|
||||
assert_eq!(upstream_url, "https://tiger.bookapi.cc/codex/responses");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_metadata_for_codex_openai_cli_requests() {
|
||||
let body_json = json!({
|
||||
|
||||
+1
-2
@@ -310,9 +310,8 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
|
||||
"upstream_url": upstream_url,
|
||||
"provider_request_method": serde_json::Value::Null,
|
||||
"provider_request_headers": provider_request_headers,
|
||||
"provider_request_body": provider_request_body,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
|
||||
"has_envelope": false,
|
||||
"needs_conversion": true,
|
||||
}),
|
||||
|
||||
+1
-2
@@ -254,9 +254,8 @@ pub(super) async fn build_same_format_local_openai_chat_decision_payload_for_can
|
||||
"upstream_url": upstream_url,
|
||||
"provider_request_method": serde_json::Value::Null,
|
||||
"provider_request_headers": provider_request_headers,
|
||||
"provider_request_body": provider_request_body,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
|
||||
"has_envelope": false,
|
||||
"needs_conversion": false,
|
||||
}),
|
||||
|
||||
+28
-3
@@ -1,6 +1,7 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::json;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::ai_pipeline::collect_control_headers;
|
||||
use crate::ai_pipeline::transport::{
|
||||
@@ -53,6 +54,31 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
|
||||
let proxy =
|
||||
resolve_transport_proxy_snapshot_with_tunnel_affinity(state, &resolved.transport).await;
|
||||
let tls_profile = resolve_transport_tls_profile(&resolved.transport);
|
||||
let timeouts = resolve_transport_execution_timeouts(&resolved.transport);
|
||||
|
||||
debug!(
|
||||
event_name = "local_openai_cli_decision_payload_built",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id,
|
||||
candidate_id = %candidate_id,
|
||||
candidate_index,
|
||||
provider_name = %resolved.transport.provider.name,
|
||||
provider_id = %candidate.provider_id,
|
||||
endpoint_id = %candidate.endpoint_id,
|
||||
key_id = %candidate.key_id,
|
||||
decision_kind = spec.decision_kind,
|
||||
execution_strategy = resolved.execution_strategy.as_str(),
|
||||
conversion_mode = resolved.conversion_mode.as_str(),
|
||||
client_api_format = spec.api_format,
|
||||
provider_api_format = %resolved.provider_api_format,
|
||||
request_path = %parts.uri.path(),
|
||||
request_query = ?parts.uri.query(),
|
||||
upstream_base_url = %resolved.transport.endpoint.base_url,
|
||||
upstream_url = %resolved.upstream_url,
|
||||
upstream_is_stream = resolved.upstream_is_stream,
|
||||
has_envelope = resolved.is_antigravity,
|
||||
"gateway built local openai cli decision payload"
|
||||
);
|
||||
|
||||
Some(GatewayControlSyncDecisionResponse {
|
||||
action: if spec.require_streaming {
|
||||
@@ -89,7 +115,7 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
tls_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&resolved.transport),
|
||||
timeouts,
|
||||
upstream_is_stream: resolved.upstream_is_stream,
|
||||
report_kind: Some(spec.report_kind.to_string()),
|
||||
report_context: Some(append_local_failover_policy_to_value(
|
||||
@@ -115,9 +141,8 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
|
||||
"upstream_url": resolved.upstream_url,
|
||||
"provider_request_method": serde_json::Value::Null,
|
||||
"provider_request_headers": resolved.provider_request_headers,
|
||||
"provider_request_body": resolved.provider_request_body,
|
||||
"original_headers": collect_control_headers(&parts.headers),
|
||||
"original_request_body": body_json,
|
||||
"original_request_body": crate::ai_pipeline::build_report_context_original_request_echo(body_json),
|
||||
"has_envelope": resolved.is_antigravity,
|
||||
"envelope_name": if resolved.is_antigravity {
|
||||
Some("antigravity:v1internal")
|
||||
|
||||
+25
-1
@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
|
||||
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::Value;
|
||||
use tracing::warn;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use crate::ai_pipeline::conversion::{
|
||||
request_conversion_direct_auth, request_conversion_kind,
|
||||
@@ -430,6 +430,30 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
|
||||
ConversionMode::None
|
||||
};
|
||||
|
||||
debug!(
|
||||
event_name = "local_openai_cli_upstream_url_resolved",
|
||||
log_type = "debug",
|
||||
trace_id = %trace_id,
|
||||
candidate_id = %candidate_id,
|
||||
candidate_index,
|
||||
provider_id = %candidate.provider_id,
|
||||
endpoint_id = %candidate.endpoint_id,
|
||||
key_id = %candidate.key_id,
|
||||
provider_type = %transport.provider.provider_type,
|
||||
client_api_format = spec.api_format,
|
||||
provider_api_format = %provider_api_format,
|
||||
execution_strategy = execution_strategy.as_str(),
|
||||
conversion_mode = conversion_mode.as_str(),
|
||||
base_url = %transport.endpoint.base_url,
|
||||
custom_path = ?transport.endpoint.custom_path,
|
||||
request_path = %parts.uri.path(),
|
||||
request_query = ?parts.uri.query(),
|
||||
mapped_model = %mapped_model,
|
||||
upstream_url = %upstream_url,
|
||||
upstream_is_stream,
|
||||
"gateway resolved local openai cli upstream url"
|
||||
);
|
||||
|
||||
Some(LocalOpenAiCliCandidatePayloadParts {
|
||||
auth_header,
|
||||
auth_value,
|
||||
|
||||
+36
-9
@@ -1,4 +1,5 @@
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
use tracing::debug;
|
||||
|
||||
use super::super::{
|
||||
augment_sync_report_context, generic_decision_missing_exact_provider_request,
|
||||
@@ -263,12 +264,12 @@ pub(crate) fn build_openai_cli_stream_plan_from_decision(
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let url = if let Some(upstream_url) = payload
|
||||
let (url, url_source) = if let Some(upstream_url) = payload
|
||||
.upstream_url
|
||||
.clone()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
{
|
||||
upstream_url
|
||||
(upstream_url, "upstream_url")
|
||||
} else {
|
||||
let Some(upstream_base_url) = payload
|
||||
.upstream_base_url
|
||||
@@ -277,7 +278,10 @@ pub(crate) fn build_openai_cli_stream_plan_from_decision(
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
build_openai_cli_url(&upstream_base_url, parts.uri.query(), compact)
|
||||
(
|
||||
build_openai_cli_url(&upstream_base_url, parts.uri.query(), compact),
|
||||
"upstream_base_url",
|
||||
)
|
||||
};
|
||||
let Some(provider_request_body_value) = payload.provider_request_body.clone() else {
|
||||
return Ok(None);
|
||||
@@ -330,6 +334,27 @@ pub(crate) fn build_openai_cli_stream_plan_from_decision(
|
||||
timeouts: payload.timeouts.clone(),
|
||||
};
|
||||
|
||||
debug!(
|
||||
event_name = "local_openai_cli_stream_plan_built",
|
||||
log_type = "debug",
|
||||
request_id = %plan.request_id,
|
||||
candidate_id = ?plan.candidate_id,
|
||||
provider_id = %plan.provider_id,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
downstream_path = %parts.uri.path(),
|
||||
downstream_query = ?parts.uri.query(),
|
||||
url_source,
|
||||
decision_upstream_base_url = ?payload.upstream_base_url,
|
||||
decision_upstream_url = ?payload.upstream_url,
|
||||
plan_url = %plan.url,
|
||||
client_api_format = %plan.client_api_format,
|
||||
provider_api_format = %plan.provider_api_format,
|
||||
upstream_is_stream = payload.upstream_is_stream,
|
||||
compact,
|
||||
"gateway built local openai cli stream execution plan"
|
||||
);
|
||||
|
||||
Ok(Some(LocalStreamPlanAndReport {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
@@ -446,12 +471,14 @@ mod tests {
|
||||
"prompt_cache_key",
|
||||
]
|
||||
);
|
||||
let report_context = built
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_request_body"))
|
||||
.expect("report context should contain provider request body");
|
||||
assert_eq!(object_keys(report_context), object_keys(plan_body));
|
||||
assert!(
|
||||
built
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_request_body"))
|
||||
.is_none(),
|
||||
"report context should not duplicate provider request body"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
use tracing::debug;
|
||||
|
||||
use super::super::{
|
||||
augment_sync_report_context, generic_decision_missing_exact_provider_request,
|
||||
@@ -267,12 +268,12 @@ pub(crate) fn build_openai_cli_sync_plan_from_decision(
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let url = if let Some(upstream_url) = payload
|
||||
let (url, url_source) = if let Some(upstream_url) = payload
|
||||
.upstream_url
|
||||
.clone()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
{
|
||||
upstream_url
|
||||
(upstream_url, "upstream_url")
|
||||
} else {
|
||||
let Some(upstream_base_url) = payload
|
||||
.upstream_base_url
|
||||
@@ -281,7 +282,10 @@ pub(crate) fn build_openai_cli_sync_plan_from_decision(
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
build_openai_cli_url(&upstream_base_url, parts.uri.query(), compact)
|
||||
(
|
||||
build_openai_cli_url(&upstream_base_url, parts.uri.query(), compact),
|
||||
"upstream_base_url",
|
||||
)
|
||||
};
|
||||
let Some(provider_request_body_value) = payload.provider_request_body.clone() else {
|
||||
return Ok(None);
|
||||
@@ -324,6 +328,27 @@ pub(crate) fn build_openai_cli_sync_plan_from_decision(
|
||||
timeouts: payload.timeouts.clone(),
|
||||
};
|
||||
|
||||
debug!(
|
||||
event_name = "local_openai_cli_sync_plan_built",
|
||||
log_type = "debug",
|
||||
request_id = %plan.request_id,
|
||||
candidate_id = ?plan.candidate_id,
|
||||
provider_id = %plan.provider_id,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
downstream_path = %parts.uri.path(),
|
||||
downstream_query = ?parts.uri.query(),
|
||||
url_source,
|
||||
decision_upstream_base_url = ?payload.upstream_base_url,
|
||||
decision_upstream_url = ?payload.upstream_url,
|
||||
plan_url = %plan.url,
|
||||
client_api_format = %plan.client_api_format,
|
||||
provider_api_format = %plan.provider_api_format,
|
||||
upstream_is_stream = payload.upstream_is_stream,
|
||||
compact,
|
||||
"gateway built local openai cli sync execution plan"
|
||||
);
|
||||
|
||||
Ok(Some(LocalSyncPlanAndReport {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
@@ -440,12 +465,14 @@ mod tests {
|
||||
"prompt_cache_key",
|
||||
]
|
||||
);
|
||||
let report_context = built
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_request_body"))
|
||||
.expect("report context should contain provider request body");
|
||||
assert_eq!(object_keys(report_context), object_keys(plan_body));
|
||||
assert!(
|
||||
built
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_request_body"))
|
||||
.is_none(),
|
||||
"report context should not duplicate provider request body"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -721,6 +721,88 @@ RETURNING
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const UPDATE_AUTH_API_KEY_WALLET_LIMIT_MODE_SQL: &str = r#"
|
||||
UPDATE wallets
|
||||
SET
|
||||
limit_mode = $2,
|
||||
updated_at = NOW()
|
||||
WHERE api_key_id = $1
|
||||
RETURNING
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
CAST(total_recharged AS DOUBLE PRECISION) AS total_recharged,
|
||||
CAST(total_consumed AS DOUBLE PRECISION) AS total_consumed,
|
||||
CAST(total_refunded AS DOUBLE PRECISION) AS total_refunded,
|
||||
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const UPDATE_AUTH_USER_WALLET_SNAPSHOT_SQL: &str = r#"
|
||||
UPDATE wallets
|
||||
SET
|
||||
balance = $2,
|
||||
gift_balance = $3,
|
||||
limit_mode = $4,
|
||||
currency = $5,
|
||||
status = $6,
|
||||
total_recharged = $7,
|
||||
total_consumed = $8,
|
||||
total_refunded = $9,
|
||||
total_adjusted = $10,
|
||||
updated_at = COALESCE(TO_TIMESTAMP($11::DOUBLE PRECISION), updated_at)
|
||||
WHERE user_id = $1
|
||||
RETURNING
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
CAST(total_recharged AS DOUBLE PRECISION) AS total_recharged,
|
||||
CAST(total_consumed AS DOUBLE PRECISION) AS total_consumed,
|
||||
CAST(total_refunded AS DOUBLE PRECISION) AS total_refunded,
|
||||
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const UPDATE_AUTH_API_KEY_WALLET_SNAPSHOT_SQL: &str = r#"
|
||||
UPDATE wallets
|
||||
SET
|
||||
balance = $2,
|
||||
gift_balance = $3,
|
||||
limit_mode = $4,
|
||||
currency = $5,
|
||||
status = $6,
|
||||
total_recharged = $7,
|
||||
total_consumed = $8,
|
||||
total_refunded = $9,
|
||||
total_adjusted = $10,
|
||||
updated_at = COALESCE(TO_TIMESTAMP($11::DOUBLE PRECISION), updated_at)
|
||||
WHERE api_key_id = $1
|
||||
RETURNING
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
CAST(total_recharged AS DOUBLE PRECISION) AS total_recharged,
|
||||
CAST(total_consumed AS DOUBLE PRECISION) AS total_consumed,
|
||||
CAST(total_refunded AS DOUBLE PRECISION) AS total_refunded,
|
||||
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const CREATE_AUTH_USER_WALLET_SQL: &str = r#"
|
||||
INSERT INTO wallets (
|
||||
id,
|
||||
@@ -770,6 +852,55 @@ RETURNING
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const CREATE_AUTH_API_KEY_WALLET_SQL: &str = r#"
|
||||
INSERT INTO wallets (
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
balance,
|
||||
gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
total_recharged,
|
||||
total_consumed,
|
||||
total_refunded,
|
||||
total_adjusted,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES (
|
||||
$1,
|
||||
NULL,
|
||||
$2,
|
||||
0,
|
||||
$3,
|
||||
$4,
|
||||
'USD',
|
||||
'active',
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
$5,
|
||||
NOW(),
|
||||
NOW()
|
||||
)
|
||||
RETURNING
|
||||
id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
CAST(total_recharged AS DOUBLE PRECISION) AS total_recharged,
|
||||
CAST(total_consumed AS DOUBLE PRECISION) AS total_consumed,
|
||||
CAST(total_refunded AS DOUBLE PRECISION) AS total_refunded,
|
||||
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
const CREATE_AUTH_USER_WALLET_GIFT_TX_SQL: &str = r#"
|
||||
INSERT INTO wallet_transactions (
|
||||
id,
|
||||
@@ -809,6 +940,45 @@ VALUES (
|
||||
)
|
||||
"#;
|
||||
|
||||
const CREATE_AUTH_API_KEY_WALLET_GIFT_TX_SQL: &str = r#"
|
||||
INSERT INTO wallet_transactions (
|
||||
id,
|
||||
wallet_id,
|
||||
category,
|
||||
reason_code,
|
||||
amount,
|
||||
balance_before,
|
||||
balance_after,
|
||||
recharge_balance_before,
|
||||
recharge_balance_after,
|
||||
gift_balance_before,
|
||||
gift_balance_after,
|
||||
link_type,
|
||||
link_id,
|
||||
operator_id,
|
||||
description,
|
||||
created_at
|
||||
)
|
||||
VALUES (
|
||||
$1,
|
||||
$2,
|
||||
'gift',
|
||||
'gift_initial',
|
||||
$3,
|
||||
0,
|
||||
$3,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
$3,
|
||||
'system_task',
|
||||
$4,
|
||||
NULL,
|
||||
'独立余额 Key 初始赠款',
|
||||
NOW()
|
||||
)
|
||||
"#;
|
||||
|
||||
fn map_user_session_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<StoredUserSessionRecord, DataLayerError> {
|
||||
@@ -1488,6 +1658,45 @@ impl GatewayDataState {
|
||||
Ok(Some(wallet))
|
||||
}
|
||||
|
||||
pub(crate) async fn initialize_auth_api_key_wallet(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
initial_gift_usd: f64,
|
||||
unlimited: bool,
|
||||
) -> Result<Option<StoredWalletSnapshot>, DataLayerError> {
|
||||
let Some(pool) = self.postgres_pool() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut tx = pool.begin().await.map_postgres_err()?;
|
||||
let gift_amount = if unlimited {
|
||||
0.0
|
||||
} else {
|
||||
initial_gift_usd.max(0.0)
|
||||
};
|
||||
let row = sqlx::query(CREATE_AUTH_API_KEY_WALLET_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(api_key_id)
|
||||
.bind(gift_amount)
|
||||
.bind(if unlimited { "unlimited" } else { "finite" })
|
||||
.bind(gift_amount)
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
let wallet = map_wallet_snapshot_row(&row)?;
|
||||
if gift_amount > 0.0 {
|
||||
sqlx::query(CREATE_AUTH_API_KEY_WALLET_GIFT_TX_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(&wallet.id)
|
||||
.bind(gift_amount)
|
||||
.bind(api_key_id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
}
|
||||
tx.commit().await.map_err(postgres_error)?;
|
||||
Ok(Some(wallet))
|
||||
}
|
||||
|
||||
pub(crate) async fn update_auth_user_wallet_limit_mode(
|
||||
&self,
|
||||
user_id: &str,
|
||||
@@ -1505,6 +1714,95 @@ impl GatewayDataState {
|
||||
row.as_ref().map(map_wallet_snapshot_row).transpose()
|
||||
}
|
||||
|
||||
pub(crate) async fn update_auth_api_key_wallet_limit_mode(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
limit_mode: &str,
|
||||
) -> Result<Option<StoredWalletSnapshot>, DataLayerError> {
|
||||
let Some(pool) = self.postgres_pool() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let row = sqlx::query(UPDATE_AUTH_API_KEY_WALLET_LIMIT_MODE_SQL)
|
||||
.bind(api_key_id)
|
||||
.bind(limit_mode)
|
||||
.fetch_optional(&pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
row.as_ref().map(map_wallet_snapshot_row).transpose()
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn update_auth_user_wallet_snapshot(
|
||||
&self,
|
||||
user_id: &str,
|
||||
balance: f64,
|
||||
gift_balance: f64,
|
||||
limit_mode: &str,
|
||||
currency: &str,
|
||||
status: &str,
|
||||
total_recharged: f64,
|
||||
total_consumed: f64,
|
||||
total_refunded: f64,
|
||||
total_adjusted: f64,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Result<Option<StoredWalletSnapshot>, DataLayerError> {
|
||||
let Some(pool) = self.postgres_pool() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let row = sqlx::query(UPDATE_AUTH_USER_WALLET_SNAPSHOT_SQL)
|
||||
.bind(user_id)
|
||||
.bind(balance)
|
||||
.bind(gift_balance)
|
||||
.bind(limit_mode)
|
||||
.bind(currency)
|
||||
.bind(status)
|
||||
.bind(total_recharged)
|
||||
.bind(total_consumed)
|
||||
.bind(total_refunded)
|
||||
.bind(total_adjusted)
|
||||
.bind(updated_at_unix_secs.map(|value| value as i64))
|
||||
.fetch_optional(&pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
row.as_ref().map(map_wallet_snapshot_row).transpose()
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn update_auth_api_key_wallet_snapshot(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
balance: f64,
|
||||
gift_balance: f64,
|
||||
limit_mode: &str,
|
||||
currency: &str,
|
||||
status: &str,
|
||||
total_recharged: f64,
|
||||
total_consumed: f64,
|
||||
total_refunded: f64,
|
||||
total_adjusted: f64,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Result<Option<StoredWalletSnapshot>, DataLayerError> {
|
||||
let Some(pool) = self.postgres_pool() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let row = sqlx::query(UPDATE_AUTH_API_KEY_WALLET_SNAPSHOT_SQL)
|
||||
.bind(api_key_id)
|
||||
.bind(balance)
|
||||
.bind(gift_balance)
|
||||
.bind(limit_mode)
|
||||
.bind(currency)
|
||||
.bind(status)
|
||||
.bind(total_recharged)
|
||||
.bind(total_consumed)
|
||||
.bind(total_refunded)
|
||||
.bind(total_adjusted)
|
||||
.bind(updated_at_unix_secs.map(|value| value as i64))
|
||||
.fetch_optional(&pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
row.as_ref().map(map_wallet_snapshot_row).transpose()
|
||||
}
|
||||
|
||||
pub(crate) async fn count_active_admin_users(&self) -> Result<u64, DataLayerError> {
|
||||
let Some(pool) = self.postgres_pool() else {
|
||||
return Ok(0);
|
||||
|
||||
@@ -639,6 +639,16 @@ impl GatewayDataState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_request_usage_body_ref(
|
||||
&self,
|
||||
body_ref: &str,
|
||||
) -> Result<Option<serde_json::Value>, DataLayerError> {
|
||||
match &self.usage_reader {
|
||||
Some(repository) => repository.resolve_body_ref(body_ref).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
|
||||
@@ -6,6 +6,10 @@ use aether_data_contracts::repository::candidates::RequestCandidateStatus;
|
||||
use aether_scheduler_core::{
|
||||
parse_request_candidate_report_context, SchedulerRequestCandidateStatusUpdate,
|
||||
};
|
||||
use aether_usage_runtime::{
|
||||
build_lifecycle_usage_seed, build_stream_terminal_usage_payload_seed,
|
||||
build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed,
|
||||
};
|
||||
use async_stream::stream;
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::http::Response;
|
||||
@@ -62,6 +66,36 @@ use crate::usage::submit_stream_report;
|
||||
use crate::usage::{GatewayStreamReportRequest, GatewaySyncReportRequest};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
fn record_sync_terminal_usage(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) {
|
||||
let context_seed = build_terminal_usage_context_seed(plan, report_context);
|
||||
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
|
||||
state
|
||||
.usage_runtime
|
||||
.record_sync_terminal(state.data.as_ref(), &context_seed, &payload_seed);
|
||||
}
|
||||
|
||||
fn record_stream_terminal_usage(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
payload: &GatewayStreamReportRequest,
|
||||
cancelled: bool,
|
||||
) {
|
||||
let context_seed = build_terminal_usage_context_seed(plan, report_context);
|
||||
let payload_seed = build_stream_terminal_usage_payload_seed(payload);
|
||||
state.usage_runtime.record_stream_terminal(
|
||||
state.data.as_ref(),
|
||||
&context_seed,
|
||||
&payload_seed,
|
||||
cancelled,
|
||||
);
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
|
||||
pub(crate) async fn execute_execution_runtime_stream(
|
||||
state: &AppState,
|
||||
@@ -73,10 +107,10 @@ pub(crate) async fn execute_execution_runtime_stream(
|
||||
mut report_context: Option<serde_json::Value>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await;
|
||||
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
|
||||
state
|
||||
.usage_runtime
|
||||
.record_pending(state.data.as_ref(), &plan, report_context.as_ref())
|
||||
.await;
|
||||
.record_pending(state.data.as_ref(), &lifecycle_seed);
|
||||
let plan_request_id_for_log = short_request_id(plan.request_id.as_str());
|
||||
let provider_name = plan.provider_name.as_deref().unwrap_or("-");
|
||||
let endpoint_id = plan.endpoint_id.as_str();
|
||||
@@ -348,6 +382,7 @@ async fn execute_stream_from_frame_stream(
|
||||
let candidate_id = plan.candidate_id.as_deref();
|
||||
let provider_name = plan.provider_name.as_deref().unwrap_or("-");
|
||||
let model_name = plan.model_name.as_deref().unwrap_or("-");
|
||||
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
|
||||
let candidate_index = parse_request_candidate_report_context(report_context.as_ref())
|
||||
.and_then(|context| context.candidate_index)
|
||||
.map(|value| value.to_string())
|
||||
@@ -527,15 +562,7 @@ async fn execute_stream_from_frame_stream(
|
||||
body_base64: body_base64.clone(),
|
||||
telemetry: None,
|
||||
};
|
||||
state
|
||||
.usage_runtime
|
||||
.record_sync_terminal(
|
||||
state.data.as_ref(),
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
&usage_payload,
|
||||
)
|
||||
.await;
|
||||
record_sync_terminal_usage(state, &plan, report_context.as_ref(), &usage_payload);
|
||||
let terminal_unix_secs = current_request_candidate_unix_ms();
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
@@ -739,15 +766,12 @@ async fn execute_stream_from_frame_stream(
|
||||
body_base64: None,
|
||||
telemetry: prefetched_telemetry.clone(),
|
||||
};
|
||||
state
|
||||
.usage_runtime
|
||||
.record_sync_terminal(
|
||||
state.data.as_ref(),
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
&payload,
|
||||
)
|
||||
.await;
|
||||
record_sync_terminal_usage(
|
||||
state,
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
&payload,
|
||||
);
|
||||
let response = submit_local_core_error_or_sync_finalize(
|
||||
state, trace_id, decision, payload,
|
||||
)
|
||||
@@ -876,17 +900,12 @@ async fn execute_stream_from_frame_stream(
|
||||
}
|
||||
|
||||
let candidate_started_unix_secs = current_request_candidate_unix_ms();
|
||||
state
|
||||
.usage_runtime
|
||||
.record_stream_started(
|
||||
state.data.as_ref(),
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
status_code,
|
||||
&headers,
|
||||
prefetched_telemetry.as_ref(),
|
||||
)
|
||||
.await;
|
||||
state.usage_runtime.record_stream_started(
|
||||
state.data.as_ref(),
|
||||
&lifecycle_seed,
|
||||
status_code,
|
||||
prefetched_telemetry.as_ref(),
|
||||
);
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
@@ -912,6 +931,7 @@ async fn execute_stream_from_frame_stream(
|
||||
let headers_for_report = headers.clone();
|
||||
let report_kind_owned = report_kind.clone();
|
||||
let report_context_owned = report_context.clone();
|
||||
let lifecycle_seed_for_report = lifecycle_seed.clone();
|
||||
let provider_prefetched_body_for_report = provider_prefetched_body.clone();
|
||||
let prefetched_body_for_report = prefetched_body.clone();
|
||||
let prefetched_chunks_for_body = prefetched_chunks.clone();
|
||||
@@ -1066,17 +1086,12 @@ async fn execute_stream_from_frame_stream(
|
||||
);
|
||||
telemetry = Some(frame_telemetry.clone());
|
||||
if should_refresh_stream_usage {
|
||||
state_for_report
|
||||
.usage_runtime
|
||||
.record_stream_started(
|
||||
state_for_report.data.as_ref(),
|
||||
&plan_for_report,
|
||||
report_context_owned.as_ref(),
|
||||
status_code,
|
||||
&headers_for_report,
|
||||
Some(&frame_telemetry),
|
||||
)
|
||||
.await;
|
||||
state_for_report.usage_runtime.record_stream_started(
|
||||
state_for_report.data.as_ref(),
|
||||
&lifecycle_seed_for_report,
|
||||
status_code,
|
||||
Some(&frame_telemetry),
|
||||
);
|
||||
usage_stream_telemetry = Some(frame_telemetry);
|
||||
}
|
||||
}
|
||||
@@ -1230,30 +1245,25 @@ async fn execute_stream_from_frame_stream(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway skipped stream report because downstream disconnected before completion"
|
||||
);
|
||||
state_for_report
|
||||
.usage_runtime
|
||||
.record_stream_terminal(
|
||||
state_for_report.data.as_ref(),
|
||||
&plan_for_report,
|
||||
report_context_owned.as_ref(),
|
||||
&GatewayStreamReportRequest {
|
||||
trace_id: trace_id_owned.clone(),
|
||||
report_kind: report_kind_owned.clone().unwrap_or_default(),
|
||||
report_context: report_context_owned.clone(),
|
||||
status_code: 499,
|
||||
headers: headers_for_report.clone(),
|
||||
provider_body_base64: (!provider_buffered_body.is_empty()).then(|| {
|
||||
base64::engine::general_purpose::STANDARD
|
||||
.encode(&provider_buffered_body)
|
||||
}),
|
||||
client_body_base64: (!buffered_body.is_empty()).then(|| {
|
||||
base64::engine::general_purpose::STANDARD.encode(&buffered_body)
|
||||
}),
|
||||
telemetry: telemetry.clone(),
|
||||
},
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
record_stream_terminal_usage(
|
||||
&state_for_report,
|
||||
&plan_for_report,
|
||||
report_context_owned.as_ref(),
|
||||
&GatewayStreamReportRequest {
|
||||
trace_id: trace_id_owned.clone(),
|
||||
report_kind: report_kind_owned.clone().unwrap_or_default(),
|
||||
report_context: report_context_owned.clone(),
|
||||
status_code: 499,
|
||||
headers: headers_for_report.clone(),
|
||||
provider_body_base64: (!provider_buffered_body.is_empty()).then(|| {
|
||||
base64::engine::general_purpose::STANDARD.encode(&provider_buffered_body)
|
||||
}),
|
||||
client_body_base64: (!buffered_body.is_empty())
|
||||
.then(|| base64::engine::general_purpose::STANDARD.encode(&buffered_body)),
|
||||
telemetry: telemetry.clone(),
|
||||
},
|
||||
true,
|
||||
);
|
||||
record_local_request_candidate_status(
|
||||
&state_for_report,
|
||||
&plan_for_report,
|
||||
@@ -1301,16 +1311,13 @@ async fn execute_stream_from_frame_stream(
|
||||
.then(|| base64::engine::general_purpose::STANDARD.encode(&buffered_body)),
|
||||
telemetry: telemetry.clone(),
|
||||
};
|
||||
state_for_report
|
||||
.usage_runtime
|
||||
.record_stream_terminal(
|
||||
state_for_report.data.as_ref(),
|
||||
&plan_for_report,
|
||||
report_context_owned.as_ref(),
|
||||
&usage_payload,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
record_stream_terminal_usage(
|
||||
&state_for_report,
|
||||
&plan_for_report,
|
||||
report_context_owned.as_ref(),
|
||||
&usage_payload,
|
||||
false,
|
||||
);
|
||||
record_local_request_candidate_status(
|
||||
&state_for_report,
|
||||
&plan_for_report,
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
use aether_contracts::{ExecutionError, ExecutionPlan, ExecutionTelemetry};
|
||||
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
|
||||
use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate;
|
||||
use aether_usage_runtime::{
|
||||
build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use base64::Engine as _;
|
||||
@@ -119,10 +122,11 @@ async fn record_stream_sync_failure(
|
||||
failure: &StreamFailureReport,
|
||||
started_at_unix_ms: Option<u64>,
|
||||
) {
|
||||
let context_seed = build_terminal_usage_context_seed(plan, report_context);
|
||||
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
|
||||
state
|
||||
.usage_runtime
|
||||
.record_sync_terminal(state.data.as_ref(), plan, report_context, payload)
|
||||
.await;
|
||||
.record_sync_terminal(state.data.as_ref(), &context_seed, &payload_seed);
|
||||
let terminal_unix_secs = current_request_candidate_unix_ms();
|
||||
record_report_request_candidate_status(
|
||||
state,
|
||||
|
||||
@@ -6,6 +6,10 @@ use aether_scheduler_core::{
|
||||
execution_error_details, parse_request_candidate_report_context,
|
||||
SchedulerRequestCandidateStatusUpdate,
|
||||
};
|
||||
use aether_usage_runtime::{
|
||||
build_lifecycle_usage_seed, build_sync_terminal_usage_payload_seed,
|
||||
build_terminal_usage_context_seed,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use base64::Engine as _;
|
||||
@@ -55,16 +59,17 @@ struct ImplicitSyncFinalizeOutcome {
|
||||
outcome: LocalCoreSyncFinalizeOutcome,
|
||||
}
|
||||
|
||||
async fn record_sync_terminal_usage(
|
||||
fn record_sync_terminal_usage(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) {
|
||||
let context_seed = build_terminal_usage_context_seed(plan, report_context);
|
||||
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
|
||||
state
|
||||
.usage_runtime
|
||||
.record_sync_terminal(state.data.as_ref(), plan, report_context, payload)
|
||||
.await;
|
||||
.record_sync_terminal(state.data.as_ref(), &context_seed, &payload_seed);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -98,10 +103,10 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
.map(|value| value.to_string())
|
||||
.unwrap_or_else(|| "-".to_string());
|
||||
let candidate_started_unix_secs = current_request_candidate_unix_ms();
|
||||
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
|
||||
state
|
||||
.usage_runtime
|
||||
.record_pending(state.data.as_ref(), &plan, report_context.as_ref())
|
||||
.await;
|
||||
.record_pending(state.data.as_ref(), &lifecycle_seed);
|
||||
#[cfg(not(test))]
|
||||
let result = {
|
||||
match DirectSyncExecutionRuntime::new()
|
||||
@@ -349,7 +354,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
.background_report
|
||||
.as_ref()
|
||||
.unwrap_or(&implicit_finalize.payload);
|
||||
record_sync_terminal_usage(state, &plan, report_context.as_ref(), usage_payload).await;
|
||||
record_sync_terminal_usage(state, &plan, report_context.as_ref(), usage_payload);
|
||||
if let Some(report_payload) = implicit_finalize.outcome.background_report {
|
||||
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
|
||||
} else {
|
||||
@@ -386,8 +391,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
&plan,
|
||||
payload.report_context.as_ref(),
|
||||
usage_payload,
|
||||
)
|
||||
.await;
|
||||
);
|
||||
if let Some(report_payload) = outcome.background_report {
|
||||
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
|
||||
} else {
|
||||
@@ -417,8 +421,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
&plan,
|
||||
payload.report_context.as_ref(),
|
||||
&outcome.report_payload,
|
||||
)
|
||||
.await;
|
||||
);
|
||||
if let Some(snapshot) = outcome.local_task_snapshot.clone() {
|
||||
state.video_tasks.record_snapshot(snapshot.clone());
|
||||
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
|
||||
@@ -454,8 +457,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
&plan,
|
||||
payload.report_context.as_ref(),
|
||||
&usage_payload,
|
||||
)
|
||||
.await;
|
||||
);
|
||||
state
|
||||
.video_tasks
|
||||
.apply_finalize_mutation(request_path, payload.report_kind.as_str());
|
||||
@@ -505,8 +507,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
&plan,
|
||||
payload.report_context.as_ref(),
|
||||
&usage_payload,
|
||||
)
|
||||
.await;
|
||||
);
|
||||
if let Some(error_report_kind) =
|
||||
resolve_local_sync_error_background_report_kind(payload.report_kind.as_str())
|
||||
{
|
||||
@@ -530,7 +531,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
record_sync_terminal_usage(state, &plan, payload.report_context.as_ref(), &payload).await;
|
||||
record_sync_terminal_usage(state, &plan, payload.report_context.as_ref(), &payload);
|
||||
let response =
|
||||
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
@@ -540,7 +541,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
)?));
|
||||
}
|
||||
|
||||
record_sync_terminal_usage(state, &plan, report_context.as_ref(), &base_usage_payload).await;
|
||||
record_sync_terminal_usage(state, &plan, report_context.as_ref(), &base_usage_payload);
|
||||
if let Some(report_kind) = report_kind {
|
||||
let report = GatewaySyncReportRequest {
|
||||
trace_id: trace_id.to_string(),
|
||||
|
||||
@@ -16,8 +16,9 @@ pub(crate) use candidate_loop::{
|
||||
};
|
||||
pub(crate) use orchestration::*;
|
||||
pub(crate) use outcome::{
|
||||
build_local_execution_exhaustion, record_failed_usage_for_exhausted_request,
|
||||
LocalExecutionExhaustion, LocalExecutionRequestOutcome,
|
||||
build_local_execution_exhaustion, build_local_execution_runtime_miss_context,
|
||||
record_failed_usage_for_exhausted_request, record_failed_usage_for_runtime_miss_request,
|
||||
LocalExecutionExhaustion, LocalExecutionRequestOutcome, LocalExecutionRuntimeMissContext,
|
||||
};
|
||||
pub(crate) use plan_fallback::{
|
||||
maybe_execute_stream_via_plan_fallback, maybe_execute_sync_via_plan_fallback,
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::time::Instant;
|
||||
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_usage_runtime::{
|
||||
build_usage_event_data_seed, UsageEvent, UsageEventData, UsageEventType,
|
||||
};
|
||||
@@ -12,7 +16,10 @@ use axum::http::{self, Response};
|
||||
use serde_json::{json, Map, Value};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::constants::LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER;
|
||||
use crate::constants::{
|
||||
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
|
||||
};
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::state::LocalExecutionRuntimeMissDiagnostic;
|
||||
use crate::AppState;
|
||||
|
||||
@@ -34,12 +41,59 @@ pub(crate) struct LocalExecutionExhaustion {
|
||||
upstream_error_message: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub(crate) struct LocalExecutionRuntimeMissContext {
|
||||
pub(crate) auth_user_id: Option<String>,
|
||||
pub(crate) auth_api_key_id: Option<String>,
|
||||
pub(crate) auth_username: Option<String>,
|
||||
pub(crate) auth_api_key_name: Option<String>,
|
||||
candidate_contexts: Vec<RuntimeMissCandidateContext>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct RuntimeMissCandidateContext {
|
||||
candidate: StoredRequestCandidate,
|
||||
provider_name: Option<String>,
|
||||
key_name: Option<String>,
|
||||
client_api_format: Option<String>,
|
||||
provider_api_format: Option<String>,
|
||||
global_model_name: Option<String>,
|
||||
selected_provider_model_name: Option<String>,
|
||||
endpoint_url: Option<String>,
|
||||
}
|
||||
|
||||
impl LocalExecutionRequestOutcome {
|
||||
pub(crate) fn responded(response: Response<Body>) -> Self {
|
||||
Self::Responded(response)
|
||||
}
|
||||
}
|
||||
|
||||
impl LocalExecutionRuntimeMissContext {
|
||||
pub(crate) fn persisted_candidate_count(&self) -> usize {
|
||||
self.candidate_contexts.len()
|
||||
}
|
||||
|
||||
pub(crate) fn candidate_summary(&self) -> Option<String> {
|
||||
const MAX_ITEMS: usize = 5;
|
||||
|
||||
if self.candidate_contexts.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut summaries = self
|
||||
.candidate_contexts
|
||||
.iter()
|
||||
.take(MAX_ITEMS)
|
||||
.map(format_runtime_miss_candidate_summary)
|
||||
.collect::<Vec<_>>();
|
||||
let remaining = self.candidate_contexts.len().saturating_sub(MAX_ITEMS);
|
||||
if remaining > 0 {
|
||||
summaries.push(format!("+{remaining} more"));
|
||||
}
|
||||
Some(summaries.join(" | "))
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_execution_exhaustion(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
@@ -100,6 +154,22 @@ pub(crate) async fn build_local_execution_exhaustion(
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_execution_runtime_miss_context(
|
||||
state: &AppState,
|
||||
request_id: &str,
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
) -> LocalExecutionRuntimeMissContext {
|
||||
let auth_context = decision.and_then(|value| value.auth_context.as_ref());
|
||||
|
||||
LocalExecutionRuntimeMissContext {
|
||||
auth_user_id: auth_context.map(|value| value.user_id.clone()),
|
||||
auth_api_key_id: auth_context.map(|value| value.api_key_id.clone()),
|
||||
auth_username: auth_context.and_then(|value| value.username.clone()),
|
||||
auth_api_key_name: auth_context.and_then(|value| value.api_key_name.clone()),
|
||||
candidate_contexts: load_runtime_miss_candidate_contexts(state, request_id, decision).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn record_failed_usage_for_exhausted_request(
|
||||
state: &AppState,
|
||||
exhaustion: LocalExecutionExhaustion,
|
||||
@@ -171,24 +241,150 @@ pub(crate) async fn record_failed_usage_for_exhausted_request(
|
||||
None => Map::new(),
|
||||
};
|
||||
request_metadata.insert("trace_id".to_string(), Value::String(request_id.clone()));
|
||||
if let Some(candidate_id) = candidate_id {
|
||||
request_metadata.insert("candidate_id".to_string(), Value::String(candidate_id));
|
||||
}
|
||||
if let Some(candidate_index) = candidate_index {
|
||||
request_metadata.insert(
|
||||
"candidate_index".to_string(),
|
||||
Value::Number(candidate_index.into()),
|
||||
);
|
||||
}
|
||||
apply_runtime_miss_usage_routing(
|
||||
&mut data,
|
||||
&mut request_metadata,
|
||||
candidate_id.as_deref(),
|
||||
candidate_index,
|
||||
None,
|
||||
diagnostic,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
data.request_metadata = Some(Value::Object(request_metadata));
|
||||
|
||||
state
|
||||
.usage_runtime
|
||||
.record_terminal_event(
|
||||
state.data.as_ref(),
|
||||
UsageEvent::new(UsageEventType::Failed, request_id, data),
|
||||
)
|
||||
.await;
|
||||
state.usage_runtime.submit_terminal_event(
|
||||
state.data.as_ref(),
|
||||
UsageEvent::new(UsageEventType::Failed, request_id, data),
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) async fn record_failed_usage_for_runtime_miss_request(
|
||||
state: &AppState,
|
||||
request_id: &str,
|
||||
started_at: &Instant,
|
||||
local_execution_runtime_miss_detail: &str,
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
diagnostic: Option<&LocalExecutionRuntimeMissDiagnostic>,
|
||||
context: &LocalExecutionRuntimeMissContext,
|
||||
) {
|
||||
if !state.usage_runtime.is_enabled() {
|
||||
return;
|
||||
}
|
||||
|
||||
let selected_candidate = select_last_runtime_miss_candidate(&context.candidate_contexts);
|
||||
let api_format = selected_candidate
|
||||
.and_then(|value| value.client_api_format.clone())
|
||||
.or_else(|| {
|
||||
trimmed_non_empty(decision.and_then(|value| value.auth_endpoint_signature.as_deref()))
|
||||
});
|
||||
let provider_api_format = selected_candidate
|
||||
.and_then(|value| value.provider_api_format.clone())
|
||||
.or_else(|| api_format.clone());
|
||||
let provider_name = selected_candidate
|
||||
.and_then(|value| value.provider_name.clone())
|
||||
.or_else(|| selected_candidate.and_then(|value| value.candidate.provider_id.clone()))
|
||||
.or_else(|| trimmed_non_empty(decision.and_then(|value| value.route_family.as_deref())))
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
let model = trimmed_non_empty(diagnostic.and_then(|value| value.requested_model.as_deref()))
|
||||
.or_else(|| selected_candidate.and_then(|value| value.global_model_name.clone()))
|
||||
.or_else(|| selected_candidate.and_then(|value| value.selected_provider_model_name.clone()))
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
let target_model = selected_candidate
|
||||
.and_then(|value| value.selected_provider_model_name.clone())
|
||||
.filter(|value| !value.eq_ignore_ascii_case(model.as_str()));
|
||||
|
||||
let status_code = http::StatusCode::SERVICE_UNAVAILABLE.as_u16();
|
||||
let client_body = json!({
|
||||
"error": {
|
||||
"type": "http_error",
|
||||
"message": local_execution_runtime_miss_detail,
|
||||
}
|
||||
});
|
||||
let mut client_headers = Map::from_iter([(
|
||||
"content-type".to_string(),
|
||||
Value::String("application/json".to_string()),
|
||||
)]);
|
||||
if let Some(reason) = diagnostic
|
||||
.map(|value| value.reason.trim())
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
client_headers.insert(
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER.to_string(),
|
||||
Value::String(reason.to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
let mut request_metadata = Map::new();
|
||||
request_metadata.insert(
|
||||
"trace_id".to_string(),
|
||||
Value::String(request_id.to_string()),
|
||||
);
|
||||
let mut data = UsageEventData {
|
||||
user_id: context.auth_user_id.clone(),
|
||||
api_key_id: context.auth_api_key_id.clone(),
|
||||
username: context.auth_username.clone(),
|
||||
api_key_name: context.auth_api_key_name.clone(),
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id: selected_candidate.and_then(|value| value.candidate.provider_id.clone()),
|
||||
provider_endpoint_id: selected_candidate
|
||||
.and_then(|value| value.candidate.endpoint_id.clone()),
|
||||
provider_api_key_id: selected_candidate.and_then(|value| value.candidate.key_id.clone()),
|
||||
request_type: Some(infer_request_type(api_format.as_deref())),
|
||||
api_format: api_format.clone(),
|
||||
api_family: api_format
|
||||
.as_deref()
|
||||
.and_then(infer_api_family)
|
||||
.map(ToOwned::to_owned),
|
||||
endpoint_kind: api_format
|
||||
.as_deref()
|
||||
.and_then(infer_endpoint_kind)
|
||||
.map(ToOwned::to_owned),
|
||||
endpoint_api_format: provider_api_format.clone(),
|
||||
provider_api_family: provider_api_format
|
||||
.as_deref()
|
||||
.and_then(infer_api_family)
|
||||
.map(ToOwned::to_owned),
|
||||
provider_endpoint_kind: provider_api_format
|
||||
.as_deref()
|
||||
.and_then(infer_endpoint_kind)
|
||||
.map(ToOwned::to_owned),
|
||||
has_format_conversion: selected_candidate.and_then(|value| {
|
||||
value
|
||||
.client_api_format
|
||||
.as_deref()
|
||||
.zip(value.provider_api_format.as_deref())
|
||||
.map(|(left, right)| !left.eq_ignore_ascii_case(right))
|
||||
}),
|
||||
status_code: Some(status_code),
|
||||
error_message: Some(local_execution_runtime_miss_detail.to_string()),
|
||||
error_category: error_category_for_failed_status(status_code),
|
||||
response_time_ms: Some(started_at.elapsed().as_millis() as u64),
|
||||
response_headers: Some(json_header_map()),
|
||||
response_body: Some(client_body.clone()),
|
||||
client_response_headers: Some(Value::Object(client_headers)),
|
||||
client_response_body: Some(client_body),
|
||||
..UsageEventData::default()
|
||||
};
|
||||
apply_runtime_miss_usage_routing(
|
||||
&mut data,
|
||||
&mut request_metadata,
|
||||
selected_candidate.map(|value| value.candidate.id.as_str()),
|
||||
selected_candidate.map(|value| value.candidate.candidate_index),
|
||||
selected_candidate.and_then(|value| value.key_name.as_deref()),
|
||||
diagnostic,
|
||||
decision.and_then(|value| value.route_family.as_deref()),
|
||||
decision.and_then(|value| value.route_kind.as_deref()),
|
||||
);
|
||||
data.request_metadata =
|
||||
(!request_metadata.is_empty()).then_some(Value::Object(request_metadata));
|
||||
|
||||
state.usage_runtime.submit_terminal_event(
|
||||
state.data.as_ref(),
|
||||
UsageEvent::new(UsageEventType::Failed, request_id, data),
|
||||
);
|
||||
}
|
||||
|
||||
fn select_last_failed_request_candidate(
|
||||
@@ -214,6 +410,22 @@ fn select_last_failed_request_candidate(
|
||||
})
|
||||
}
|
||||
|
||||
fn select_last_runtime_miss_candidate(
|
||||
candidates: &[RuntimeMissCandidateContext],
|
||||
) -> Option<&RuntimeMissCandidateContext> {
|
||||
candidates.iter().max_by_key(|candidate| {
|
||||
(
|
||||
candidate.candidate.retry_index,
|
||||
candidate.candidate.candidate_index,
|
||||
candidate
|
||||
.candidate
|
||||
.finished_at_unix_ms
|
||||
.or(candidate.candidate.started_at_unix_ms)
|
||||
.unwrap_or(candidate.candidate.created_at_unix_ms),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn error_category_for_failed_status(status_code: u16) -> Option<String> {
|
||||
if status_code >= 500 {
|
||||
Some("server_error".to_string())
|
||||
@@ -230,3 +442,408 @@ fn json_header_map() -> Value {
|
||||
Value::String("application/json".to_string()),
|
||||
)]))
|
||||
}
|
||||
|
||||
async fn load_runtime_miss_candidate_contexts(
|
||||
state: &AppState,
|
||||
request_id: &str,
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
) -> Vec<RuntimeMissCandidateContext> {
|
||||
let mut candidates = match state
|
||||
.read_request_candidates_by_request_id(request_id)
|
||||
.await
|
||||
{
|
||||
Ok(candidates) => candidates,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
request_id = %request_id,
|
||||
error = ?err,
|
||||
"gateway failed to load request candidates for local execution runtime miss"
|
||||
);
|
||||
return Vec::new();
|
||||
}
|
||||
};
|
||||
if candidates.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
candidates.sort_by_key(|candidate| {
|
||||
(
|
||||
candidate.candidate_index,
|
||||
candidate.retry_index,
|
||||
candidate.created_at_unix_ms,
|
||||
)
|
||||
});
|
||||
|
||||
let (providers_by_id, endpoints_by_id, keys_by_id) = if state.has_provider_catalog_data_reader()
|
||||
{
|
||||
let provider_ids = collect_present_ids(
|
||||
candidates
|
||||
.iter()
|
||||
.filter_map(|value| value.provider_id.as_deref()),
|
||||
);
|
||||
let endpoint_ids = collect_present_ids(
|
||||
candidates
|
||||
.iter()
|
||||
.filter_map(|value| value.endpoint_id.as_deref()),
|
||||
);
|
||||
let key_ids = collect_present_ids(
|
||||
candidates
|
||||
.iter()
|
||||
.filter_map(|value| value.key_id.as_deref()),
|
||||
);
|
||||
let (providers_result, endpoints_result, keys_result) = tokio::join!(
|
||||
state.read_provider_catalog_providers_by_ids(&provider_ids),
|
||||
state.read_provider_catalog_endpoints_by_ids(&endpoint_ids),
|
||||
state.read_provider_catalog_keys_by_ids(&key_ids),
|
||||
);
|
||||
(
|
||||
match providers_result {
|
||||
Ok(values) => values
|
||||
.into_iter()
|
||||
.map(|value| (value.id.clone(), value))
|
||||
.collect::<BTreeMap<_, _>>(),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
request_id = %request_id,
|
||||
error = ?err,
|
||||
"gateway failed to load provider catalog providers for local execution runtime miss"
|
||||
);
|
||||
BTreeMap::new()
|
||||
}
|
||||
},
|
||||
match endpoints_result {
|
||||
Ok(values) => values
|
||||
.into_iter()
|
||||
.map(|value| (value.id.clone(), value))
|
||||
.collect::<BTreeMap<_, _>>(),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
request_id = %request_id,
|
||||
error = ?err,
|
||||
"gateway failed to load provider catalog endpoints for local execution runtime miss"
|
||||
);
|
||||
BTreeMap::new()
|
||||
}
|
||||
},
|
||||
match keys_result {
|
||||
Ok(values) => values
|
||||
.into_iter()
|
||||
.map(|value| (value.id.clone(), value))
|
||||
.collect::<BTreeMap<_, _>>(),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
request_id = %request_id,
|
||||
error = ?err,
|
||||
"gateway failed to load provider catalog keys for local execution runtime miss"
|
||||
);
|
||||
BTreeMap::new()
|
||||
}
|
||||
},
|
||||
)
|
||||
} else {
|
||||
(BTreeMap::new(), BTreeMap::new(), BTreeMap::new())
|
||||
};
|
||||
|
||||
candidates
|
||||
.into_iter()
|
||||
.map(|candidate| {
|
||||
let provider = candidate
|
||||
.provider_id
|
||||
.as_deref()
|
||||
.and_then(|value| providers_by_id.get(value));
|
||||
let endpoint = candidate
|
||||
.endpoint_id
|
||||
.as_deref()
|
||||
.and_then(|value| endpoints_by_id.get(value));
|
||||
let key = candidate
|
||||
.key_id
|
||||
.as_deref()
|
||||
.and_then(|value| keys_by_id.get(value));
|
||||
RuntimeMissCandidateContext {
|
||||
provider_name: candidate_extra_data_string(&candidate, "provider_name")
|
||||
.or_else(|| provider.map(|value| value.name.clone())),
|
||||
key_name: candidate_extra_data_string(&candidate, "key_name")
|
||||
.or_else(|| key.map(|value| value.name.clone())),
|
||||
client_api_format: candidate_extra_data_string(&candidate, "client_api_format")
|
||||
.or_else(|| candidate_extra_data_string(&candidate, "client_contract")),
|
||||
provider_api_format: candidate_extra_data_string(&candidate, "provider_api_format")
|
||||
.or_else(|| candidate_extra_data_string(&candidate, "provider_contract"))
|
||||
.or_else(|| endpoint.map(|value| value.api_format.clone())),
|
||||
global_model_name: candidate_extra_data_string(&candidate, "global_model_name"),
|
||||
selected_provider_model_name: candidate_extra_data_string(
|
||||
&candidate,
|
||||
"selected_provider_model_name",
|
||||
),
|
||||
endpoint_url: endpoint
|
||||
.and_then(|value| build_runtime_miss_candidate_endpoint_url(value, decision)),
|
||||
candidate,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn collect_present_ids<'a>(ids: impl Iterator<Item = &'a str>) -> Vec<String> {
|
||||
ids.filter_map(|value| {
|
||||
let trimmed = value.trim();
|
||||
(!trimmed.is_empty()).then_some(trimmed.to_string())
|
||||
})
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn candidate_extra_data_string(candidate: &StoredRequestCandidate, key: &str) -> Option<String> {
|
||||
candidate
|
||||
.extra_data
|
||||
.as_ref()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|value| value.get(key))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn build_runtime_miss_candidate_endpoint_url(
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
) -> Option<String> {
|
||||
let path = endpoint
|
||||
.custom_path
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.or_else(|| {
|
||||
decision
|
||||
.map(|value| value.public_path.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
});
|
||||
let query = decision
|
||||
.and_then(|value| value.public_query_string.as_deref())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
path.and_then(|value| {
|
||||
crate::provider_transport::url::build_passthrough_path_url(
|
||||
&endpoint.base_url,
|
||||
value,
|
||||
query,
|
||||
&[],
|
||||
)
|
||||
})
|
||||
.or_else(|| trimmed_non_empty(Some(endpoint.base_url.as_str())))
|
||||
}
|
||||
|
||||
fn format_runtime_miss_candidate_summary(candidate: &RuntimeMissCandidateContext) -> String {
|
||||
let mut parts = Vec::new();
|
||||
parts.push(format!("idx={}", candidate.candidate.candidate_index));
|
||||
parts.push(format!("retry={}", candidate.candidate.retry_index));
|
||||
parts.push(format!(
|
||||
"status={}",
|
||||
request_candidate_status_label(candidate.candidate.status)
|
||||
));
|
||||
if let Some(provider_label) = format_name_with_id(
|
||||
candidate.provider_name.as_deref(),
|
||||
candidate.candidate.provider_id.as_deref(),
|
||||
) {
|
||||
parts.push(format!("provider={provider_label}"));
|
||||
}
|
||||
if let Some(endpoint_id) = candidate
|
||||
.candidate
|
||||
.endpoint_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
parts.push(format!("endpoint={endpoint_id}"));
|
||||
}
|
||||
if let Some(endpoint_url) = candidate
|
||||
.endpoint_url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
parts.push(format!("url={endpoint_url}"));
|
||||
}
|
||||
if let Some(key_label) = format_name_with_id(
|
||||
candidate.key_name.as_deref(),
|
||||
candidate.candidate.key_id.as_deref(),
|
||||
) {
|
||||
parts.push(format!("key={key_label}"));
|
||||
}
|
||||
if let Some(skip_reason) = candidate
|
||||
.candidate
|
||||
.skip_reason
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
parts.push(format!("skip={skip_reason}"));
|
||||
}
|
||||
if let Some(status_code) = candidate.candidate.status_code {
|
||||
parts.push(format!("code={status_code}"));
|
||||
}
|
||||
if let Some(error_type) = candidate
|
||||
.candidate
|
||||
.error_type
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
parts.push(format!("error_type={error_type}"));
|
||||
}
|
||||
parts.join(" ")
|
||||
}
|
||||
|
||||
fn format_name_with_id(name: Option<&str>, id: Option<&str>) -> Option<String> {
|
||||
let name = name.map(str::trim).filter(|value| !value.is_empty());
|
||||
let id = id.map(str::trim).filter(|value| !value.is_empty());
|
||||
|
||||
match (name, id) {
|
||||
(Some(name), Some(id)) => Some(format!("{name}({id})")),
|
||||
(Some(name), None) => Some(name.to_string()),
|
||||
(None, Some(id)) => Some(id.to_string()),
|
||||
(None, None) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn request_candidate_status_label(status: RequestCandidateStatus) -> &'static str {
|
||||
match status {
|
||||
RequestCandidateStatus::Available => "available",
|
||||
RequestCandidateStatus::Unused => "unused",
|
||||
RequestCandidateStatus::Pending => "pending",
|
||||
RequestCandidateStatus::Streaming => "streaming",
|
||||
RequestCandidateStatus::Success => "success",
|
||||
RequestCandidateStatus::Failed => "failed",
|
||||
RequestCandidateStatus::Cancelled => "cancelled",
|
||||
RequestCandidateStatus::Skipped => "skipped",
|
||||
}
|
||||
}
|
||||
|
||||
fn infer_request_type(api_format: Option<&str>) -> String {
|
||||
match infer_endpoint_kind(api_format.unwrap_or_default()) {
|
||||
Some("video") => "video".to_string(),
|
||||
Some("image") => "image".to_string(),
|
||||
_ => "chat".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn infer_api_family(api_format: &str) -> Option<&str> {
|
||||
api_format.split_once(':').map(|(family, _)| family)
|
||||
}
|
||||
|
||||
fn infer_endpoint_kind(api_format: &str) -> Option<&str> {
|
||||
api_format.split_once(':').map(|(_, kind)| kind)
|
||||
}
|
||||
|
||||
fn apply_runtime_miss_usage_routing(
|
||||
data: &mut UsageEventData,
|
||||
request_metadata: &mut Map<String, Value>,
|
||||
candidate_id: Option<&str>,
|
||||
candidate_index: Option<u32>,
|
||||
key_name: Option<&str>,
|
||||
diagnostic: Option<&LocalExecutionRuntimeMissDiagnostic>,
|
||||
route_family_fallback: Option<&str>,
|
||||
route_kind_fallback: Option<&str>,
|
||||
) {
|
||||
data.candidate_id = data
|
||||
.candidate_id
|
||||
.clone()
|
||||
.or_else(|| trimmed_non_empty(candidate_id));
|
||||
data.candidate_index = data
|
||||
.candidate_index
|
||||
.or_else(|| candidate_index.map(u64::from));
|
||||
data.key_name = data
|
||||
.key_name
|
||||
.clone()
|
||||
.or_else(|| trimmed_non_empty(key_name));
|
||||
data.execution_path = data
|
||||
.execution_path
|
||||
.clone()
|
||||
.or_else(|| Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS.to_string()));
|
||||
data.local_execution_runtime_miss_reason = data
|
||||
.local_execution_runtime_miss_reason
|
||||
.clone()
|
||||
.or_else(|| trimmed_non_empty(diagnostic.map(|value| value.reason.as_str())));
|
||||
data.route_family = data.route_family.clone().or_else(|| {
|
||||
trimmed_non_empty(
|
||||
diagnostic
|
||||
.and_then(|value| value.route_family.as_deref())
|
||||
.or(route_family_fallback),
|
||||
)
|
||||
});
|
||||
data.route_kind = data.route_kind.clone().or_else(|| {
|
||||
trimmed_non_empty(
|
||||
diagnostic
|
||||
.and_then(|value| value.route_kind.as_deref())
|
||||
.or(route_kind_fallback),
|
||||
)
|
||||
});
|
||||
data.planner_kind = data
|
||||
.planner_kind
|
||||
.clone()
|
||||
.or_else(|| trimmed_non_empty(diagnostic.and_then(|value| value.plan_kind.as_deref())));
|
||||
let _ = request_metadata;
|
||||
}
|
||||
|
||||
fn trimmed_non_empty(value: Option<&str>) -> Option<String> {
|
||||
value
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::apply_runtime_miss_usage_routing;
|
||||
use crate::constants::EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS;
|
||||
use crate::state::LocalExecutionRuntimeMissDiagnostic;
|
||||
use aether_usage_runtime::UsageEventData;
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
#[test]
|
||||
fn runtime_miss_routing_moves_to_typed_usage_fields_and_keeps_metadata_lightweight() {
|
||||
let mut data = UsageEventData::default();
|
||||
let mut request_metadata =
|
||||
Map::from_iter([("trace_id".to_string(), Value::String("trace-1".to_string()))]);
|
||||
|
||||
apply_runtime_miss_usage_routing(
|
||||
&mut data,
|
||||
&mut request_metadata,
|
||||
Some("cand-1"),
|
||||
Some(2),
|
||||
Some("primary"),
|
||||
Some(&LocalExecutionRuntimeMissDiagnostic {
|
||||
reason: "all_candidates_skipped".to_string(),
|
||||
route_family: Some("claude".to_string()),
|
||||
route_kind: Some("cli".to_string()),
|
||||
plan_kind: Some("claude_cli_sync".to_string()),
|
||||
..LocalExecutionRuntimeMissDiagnostic::default()
|
||||
}),
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(data.candidate_id.as_deref(), Some("cand-1"));
|
||||
assert_eq!(data.candidate_index, Some(2));
|
||||
assert_eq!(data.key_name.as_deref(), Some("primary"));
|
||||
assert_eq!(data.planner_kind.as_deref(), Some("claude_cli_sync"));
|
||||
assert_eq!(data.route_family.as_deref(), Some("claude"));
|
||||
assert_eq!(data.route_kind.as_deref(), Some("cli"));
|
||||
assert_eq!(
|
||||
data.execution_path.as_deref(),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
data.local_execution_runtime_miss_reason.as_deref(),
|
||||
Some("all_candidates_skipped")
|
||||
);
|
||||
assert_eq!(
|
||||
Value::Object(request_metadata),
|
||||
json!({
|
||||
"trace_id": "trace-1"
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -106,6 +106,12 @@ pub(super) async fn build_admin_create_api_key_response(
|
||||
allowed_models,
|
||||
rate_limit,
|
||||
concurrent_limit: 5,
|
||||
force_capabilities: None,
|
||||
is_active: true,
|
||||
expires_at_unix_secs: None,
|
||||
auto_delete_on_expiry: false,
|
||||
total_requests: 0,
|
||||
total_cost_usd: 0.0,
|
||||
},
|
||||
)
|
||||
.await?
|
||||
|
||||
@@ -32,10 +32,12 @@ pub(super) async fn load_user_leaderboard_metadata(
|
||||
},
|
||||
);
|
||||
}
|
||||
return Ok(metadata);
|
||||
}
|
||||
|
||||
for user_id in user_ids {
|
||||
if metadata.contains_key(user_id) {
|
||||
continue;
|
||||
}
|
||||
let Some(user) = state.find_user_auth_by_id(user_id).await? else {
|
||||
continue;
|
||||
};
|
||||
|
||||
@@ -130,24 +130,32 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
|
||||
let usage = state
|
||||
.list_admin_usage_for_optional_range(time_range.as_ref(), &filters)
|
||||
.await?;
|
||||
let api_key_ids: Vec<String> = usage
|
||||
.iter()
|
||||
.filter_map(|item| item.api_key_id.clone())
|
||||
.collect::<std::collections::BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect();
|
||||
let snapshots = if state.has_auth_api_key_data_reader() {
|
||||
let api_key_ids: Vec<String> = usage
|
||||
.iter()
|
||||
.filter_map(|item| item.api_key_id.clone())
|
||||
.collect::<std::collections::BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect();
|
||||
Some(
|
||||
state
|
||||
.list_auth_api_key_snapshots_by_ids(&api_key_ids)
|
||||
.resolve_auth_api_key_snapshots_by_ids(&api_key_ids)
|
||||
.await?,
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let api_key_names = if state.has_auth_api_key_data_reader() {
|
||||
state
|
||||
.resolve_auth_api_key_names_by_ids(&api_key_ids)
|
||||
.await?
|
||||
} else {
|
||||
std::collections::BTreeMap::new()
|
||||
};
|
||||
let mut leaderboard = build_api_key_leaderboard_items(
|
||||
&usage,
|
||||
snapshots.as_deref(),
|
||||
&api_key_names,
|
||||
include_inactive,
|
||||
exclude_admin,
|
||||
);
|
||||
@@ -220,8 +228,13 @@ pub(super) async fn maybe_build_local_admin_stats_leaderboard_response(
|
||||
.into_iter()
|
||||
.collect();
|
||||
let user_metadata = load_user_leaderboard_metadata(state, &user_ids).await?;
|
||||
let mut leaderboard =
|
||||
build_user_leaderboard_items(&usage, &user_metadata, include_inactive, exclude_admin);
|
||||
let mut leaderboard = build_user_leaderboard_items(
|
||||
&usage,
|
||||
&user_metadata,
|
||||
state.has_auth_user_data_reader(),
|
||||
include_inactive,
|
||||
exclude_admin,
|
||||
);
|
||||
leaderboard.sort_by(|left, right| compare_leaderboard_items(metric, order, left, right));
|
||||
|
||||
return Ok(Some(build_admin_stats_leaderboard_response(
|
||||
|
||||
@@ -28,3 +28,24 @@ pub(in super::super) async fn admin_usage_provider_key_names(
|
||||
.map(|key| (key.id, key.name))
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub(in super::super) async fn admin_usage_api_key_names(
|
||||
state: &AdminAppState<'_>,
|
||||
usage: &[StoredRequestUsageAudit],
|
||||
) -> Result<BTreeMap<String, String>, GatewayError> {
|
||||
if !state.has_auth_api_key_data_reader() {
|
||||
return Ok(BTreeMap::new());
|
||||
}
|
||||
|
||||
let api_key_ids = usage
|
||||
.iter()
|
||||
.filter_map(|item| item.api_key_id.clone())
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
if api_key_ids.is_empty() {
|
||||
return Ok(BTreeMap::new());
|
||||
}
|
||||
|
||||
state.resolve_auth_api_key_names_by_ids(&api_key_ids).await
|
||||
}
|
||||
|
||||
@@ -4,4 +4,4 @@ mod filters;
|
||||
|
||||
pub(super) use aggregations::admin_usage_aggregation_by_user_json;
|
||||
pub(super) use cache_affinity::list_recent_completed_usage_for_cache_affinity;
|
||||
pub(super) use filters::admin_usage_provider_key_names;
|
||||
pub(super) use filters::{admin_usage_api_key_names, admin_usage_provider_key_names};
|
||||
|
||||
+23
-8
@@ -16,6 +16,20 @@ use axum::{
|
||||
use serde_json::json;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
async fn load_usage_cache_affinity_usernames(
|
||||
state: &AdminAppState<'_>,
|
||||
user_ids: &[String],
|
||||
) -> Result<BTreeMap<String, String>, GatewayError> {
|
||||
Ok(state
|
||||
.resolve_auth_user_summaries_by_ids(user_ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.filter_map(|(user_id, user)| {
|
||||
(!user.username.trim().is_empty()).then_some((user_id, user.username))
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub(super) async fn build_admin_usage_cache_affinity_interval_timeline_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
@@ -41,6 +55,7 @@ pub(super) async fn build_admin_usage_cache_affinity_interval_timeline_response(
|
||||
list_recent_completed_usage_for_cache_affinity(state, hours, user_id.as_deref()).await?;
|
||||
let mut grouped: BTreeMap<String, Vec<serde_json::Value>> = BTreeMap::new();
|
||||
let mut models = BTreeSet::new();
|
||||
let mut legacy_usernames_by_user_id = BTreeMap::new();
|
||||
let mut usernames_by_user_id = BTreeMap::new();
|
||||
|
||||
for (group_user_id, items) in admin_usage_group_completed_by_user(&usage) {
|
||||
@@ -67,7 +82,7 @@ pub(super) async fn build_admin_usage_cache_affinity_interval_timeline_response(
|
||||
if include_user_info && user_id.is_none() {
|
||||
point["user_id"] = json!(group_user_id.clone());
|
||||
if let Some(username) = item.username.clone() {
|
||||
usernames_by_user_id
|
||||
legacy_usernames_by_user_id
|
||||
.entry(group_user_id.clone())
|
||||
.or_insert(username);
|
||||
}
|
||||
@@ -82,17 +97,17 @@ pub(super) async fn build_admin_usage_cache_affinity_interval_timeline_response(
|
||||
}
|
||||
}
|
||||
|
||||
if include_user_info && user_id.is_none() && state.has_user_data_reader() {
|
||||
if include_user_info && user_id.is_none() {
|
||||
let user_ids: Vec<_> = grouped.keys().cloned().collect();
|
||||
let user_map: BTreeMap<_, _> = state
|
||||
.list_users_by_ids(&user_ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|user| (user.id, user.username))
|
||||
.collect();
|
||||
let user_map = load_usage_cache_affinity_usernames(state, &user_ids).await?;
|
||||
for (user_id, username) in user_map {
|
||||
usernames_by_user_id.insert(user_id, username);
|
||||
}
|
||||
if !state.has_auth_user_data_reader() {
|
||||
for (user_id, username) in legacy_usernames_by_user_id {
|
||||
usernames_by_user_id.entry(user_id).or_insert(username);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let total_points_before_limit: usize = grouped.values().map(Vec::len).sum();
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
use super::analytics::admin_usage_provider_key_names;
|
||||
use super::analytics::{admin_usage_api_key_names, admin_usage_provider_key_names};
|
||||
use super::replay::{
|
||||
admin_usage_curl_headers, admin_usage_curl_url, admin_usage_headers_from_value,
|
||||
admin_usage_id_from_action_path, admin_usage_id_from_detail_path,
|
||||
admin_usage_resolve_request_preview_body, build_admin_usage_curl_response,
|
||||
admin_usage_resolve_body_value, admin_usage_resolve_request_preview_body,
|
||||
admin_usage_resolve_request_preview_body_for_item, build_admin_usage_curl_response,
|
||||
build_admin_usage_detail_payload, build_admin_usage_replay_response,
|
||||
};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
@@ -12,6 +13,7 @@ use aether_admin::observability::usage::{
|
||||
admin_usage_bad_request_response, admin_usage_data_unavailable_response,
|
||||
admin_usage_provider_key_name, ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
|
||||
};
|
||||
use aether_data_contracts::repository::usage::UsageBodyField;
|
||||
use axum::{
|
||||
body::Body,
|
||||
http,
|
||||
@@ -82,10 +84,22 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
|
||||
.and_then(admin_usage_headers_from_value)
|
||||
.filter(|headers| !headers.is_empty())
|
||||
.unwrap_or_else(admin_usage_curl_headers);
|
||||
let body = item
|
||||
.provider_request_body
|
||||
.clone()
|
||||
.or_else(|| item.request_body.clone())
|
||||
let provider_request_body = admin_usage_resolve_body_value(
|
||||
state,
|
||||
&item,
|
||||
item.provider_request_body.as_ref(),
|
||||
UsageBodyField::ProviderRequestBody,
|
||||
)
|
||||
.await?;
|
||||
let request_body = admin_usage_resolve_body_value(
|
||||
state,
|
||||
&item,
|
||||
item.request_body.as_ref(),
|
||||
UsageBodyField::RequestBody,
|
||||
)
|
||||
.await?;
|
||||
let body = provider_request_body
|
||||
.or(request_body)
|
||||
.unwrap_or_else(|| admin_usage_resolve_request_preview_body(&item, None));
|
||||
return Ok(Some(attach_admin_audit_response(
|
||||
build_admin_usage_curl_response(&item, url, headers_json, &headers, &body),
|
||||
@@ -146,32 +160,50 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
|
||||
};
|
||||
|
||||
let users_by_id: BTreeMap<String, aether_data::repository::users::StoredUserSummary> =
|
||||
if state.has_user_data_reader() {
|
||||
if let Some(user_id) = item.user_id.as_ref() {
|
||||
state
|
||||
.list_users_by_ids(std::slice::from_ref(user_id))
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|user| (user.id.clone(), user))
|
||||
.collect()
|
||||
} else {
|
||||
BTreeMap::new()
|
||||
}
|
||||
} else {
|
||||
BTreeMap::new()
|
||||
};
|
||||
state
|
||||
.resolve_auth_user_summaries_by_ids(
|
||||
&item.user_id.clone().into_iter().collect::<Vec<_>>(),
|
||||
)
|
||||
.await?;
|
||||
let provider_key_names =
|
||||
admin_usage_provider_key_names(state, std::slice::from_ref(&item)).await?;
|
||||
let api_key_names =
|
||||
admin_usage_api_key_names(state, std::slice::from_ref(&item)).await?;
|
||||
let provider_key_name = admin_usage_provider_key_name(&item, &provider_key_names);
|
||||
|
||||
let request_body = item
|
||||
.request_body
|
||||
.clone()
|
||||
.unwrap_or_else(|| admin_usage_resolve_request_preview_body(&item, None));
|
||||
let request_body =
|
||||
admin_usage_resolve_request_preview_body_for_item(state, &item, None).await?;
|
||||
let mut detail_item = item.clone();
|
||||
if include_bodies {
|
||||
detail_item.provider_request_body = admin_usage_resolve_body_value(
|
||||
state,
|
||||
&item,
|
||||
item.provider_request_body.as_ref(),
|
||||
UsageBodyField::ProviderRequestBody,
|
||||
)
|
||||
.await?;
|
||||
detail_item.response_body = admin_usage_resolve_body_value(
|
||||
state,
|
||||
&item,
|
||||
item.response_body.as_ref(),
|
||||
UsageBodyField::ResponseBody,
|
||||
)
|
||||
.await?;
|
||||
detail_item.client_response_body = admin_usage_resolve_body_value(
|
||||
state,
|
||||
&item,
|
||||
item.client_response_body.as_ref(),
|
||||
UsageBodyField::ClientResponseBody,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
let default_headers = admin_usage_curl_headers();
|
||||
let payload = build_admin_usage_detail_payload(
|
||||
&item,
|
||||
&detail_item,
|
||||
&users_by_id,
|
||||
&api_key_names,
|
||||
state.has_auth_user_data_reader(),
|
||||
state.has_auth_api_key_data_reader(),
|
||||
provider_key_name.as_deref(),
|
||||
include_bodies,
|
||||
request_body,
|
||||
|
||||
@@ -5,7 +5,8 @@ use aether_admin::observability::usage::{
|
||||
ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
|
||||
};
|
||||
use aether_data_contracts::repository::{
|
||||
provider_catalog::StoredProviderCatalogEndpoint, usage::StoredRequestUsageAudit,
|
||||
provider_catalog::StoredProviderCatalogEndpoint,
|
||||
usage::{StoredRequestUsageAudit, UsageBodyField},
|
||||
};
|
||||
use axum::{
|
||||
body::Body,
|
||||
@@ -46,6 +47,50 @@ pub(super) fn admin_usage_resolve_request_preview_body(
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) async fn admin_usage_resolve_body_value(
|
||||
state: &AdminAppState<'_>,
|
||||
item: &StoredRequestUsageAudit,
|
||||
inline_body: Option<&Value>,
|
||||
field: UsageBodyField,
|
||||
) -> Result<Option<Value>, GatewayError> {
|
||||
let resolved_ref_body = match item.body_ref(field) {
|
||||
Some(body_ref) => state.resolve_request_usage_body_ref(body_ref).await?,
|
||||
None => None,
|
||||
};
|
||||
Ok(admin_usage_body_value_from_sources(
|
||||
resolved_ref_body,
|
||||
inline_body,
|
||||
))
|
||||
}
|
||||
|
||||
fn admin_usage_body_value_from_sources(
|
||||
resolved_ref_body: Option<Value>,
|
||||
inline_body: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
resolved_ref_body.or_else(|| inline_body.cloned())
|
||||
}
|
||||
|
||||
pub(super) async fn admin_usage_resolve_request_preview_body_for_item(
|
||||
state: &AdminAppState<'_>,
|
||||
item: &StoredRequestUsageAudit,
|
||||
body_override: Option<serde_json::Value>,
|
||||
) -> Result<serde_json::Value, GatewayError> {
|
||||
if let Some(body_override) = body_override {
|
||||
return Ok(body_override);
|
||||
}
|
||||
if let Some(body) = admin_usage_resolve_body_value(
|
||||
state,
|
||||
item,
|
||||
item.request_body.as_ref(),
|
||||
UsageBodyField::RequestBody,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(body);
|
||||
}
|
||||
Ok(admin_usage_resolve_request_preview_body(item, None))
|
||||
}
|
||||
|
||||
pub(super) fn build_admin_usage_curl_response(
|
||||
item: &StoredRequestUsageAudit,
|
||||
url: Option<String>,
|
||||
@@ -62,9 +107,46 @@ pub(super) fn build_admin_usage_curl_response(
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::admin_usage_body_value_from_sources;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn resolved_reference_body_wins_over_inline_fallback() {
|
||||
let inline_body = json!({
|
||||
"truncated": true,
|
||||
"reason": "usage_capture_limits_exceeded"
|
||||
});
|
||||
let ref_body = json!({
|
||||
"messages": [{"role": "user", "content": "real request body"}]
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
admin_usage_body_value_from_sources(Some(ref_body.clone()), Some(&inline_body)),
|
||||
Some(ref_body)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inline_body_is_used_when_reference_body_is_unavailable() {
|
||||
let inline_body = json!({
|
||||
"messages": [{"role": "user", "content": "fallback inline body"}]
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
admin_usage_body_value_from_sources(None, Some(&inline_body)),
|
||||
Some(inline_body)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn build_admin_usage_detail_payload(
|
||||
item: &StoredRequestUsageAudit,
|
||||
users_by_id: &BTreeMap<String, aether_data::repository::users::StoredUserSummary>,
|
||||
api_key_names: &BTreeMap<String, String>,
|
||||
auth_user_reader_available: bool,
|
||||
auth_api_key_reader_available: bool,
|
||||
provider_key_name: Option<&str>,
|
||||
include_bodies: bool,
|
||||
request_body: Value,
|
||||
@@ -73,6 +155,9 @@ pub(super) fn build_admin_usage_detail_payload(
|
||||
aether_admin::observability::usage::build_admin_usage_detail_payload(
|
||||
item,
|
||||
users_by_id,
|
||||
api_key_names,
|
||||
auth_user_reader_available,
|
||||
auth_api_key_reader_available,
|
||||
provider_key_name,
|
||||
include_bodies,
|
||||
request_body,
|
||||
@@ -254,7 +339,9 @@ pub(super) async fn build_admin_usage_replay_response(
|
||||
|
||||
let same_provider = item.provider_id.as_deref() == Some(target_provider.id.as_str());
|
||||
let same_endpoint = item.provider_endpoint_id.as_deref() == Some(target_endpoint.id.as_str());
|
||||
let request_body = admin_usage_resolve_request_preview_body(&item, payload.body_override);
|
||||
let request_body =
|
||||
admin_usage_resolve_request_preview_body_for_item(state, &item, payload.body_override)
|
||||
.await?;
|
||||
|
||||
let url = admin_usage_curl_url(state, &target_endpoint, &item);
|
||||
let headers = admin_usage_curl_headers();
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use super::super::stats::{AdminStatsTimeRange, AdminStatsUsageFilter};
|
||||
use super::analytics::admin_usage_provider_key_names;
|
||||
use super::analytics::{admin_usage_api_key_names, admin_usage_provider_key_names};
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::query_param_value;
|
||||
use crate::GatewayError;
|
||||
@@ -85,10 +85,13 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
|
||||
if requested_ids.is_none() && items.len() > 50 {
|
||||
items.truncate(50);
|
||||
}
|
||||
let api_key_names = admin_usage_api_key_names(state, &items).await?;
|
||||
let provider_key_names = admin_usage_provider_key_names(state, &items).await?;
|
||||
|
||||
return Ok(Some(build_admin_usage_active_requests_response(
|
||||
&items,
|
||||
&api_key_names,
|
||||
state.has_auth_api_key_data_reader(),
|
||||
&provider_key_names,
|
||||
)));
|
||||
}
|
||||
@@ -134,10 +137,30 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
|
||||
Err(detail) => return Ok(Some(admin_usage_bad_request_response(detail))),
|
||||
};
|
||||
|
||||
let user_ids: Vec<String> = usage
|
||||
.iter()
|
||||
.filter_map(|item| item.user_id.clone())
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect();
|
||||
let users_by_id: BTreeMap<String, aether_data::repository::users::StoredUserSummary> =
|
||||
state.resolve_auth_user_summaries_by_ids(&user_ids).await?;
|
||||
let api_key_names = admin_usage_api_key_names(state, &usage).await?;
|
||||
|
||||
usage.retain(|item| {
|
||||
admin_usage_matches_search(item, search.as_deref())
|
||||
&& admin_usage_matches_username(item, username_filter.as_deref())
|
||||
&& admin_usage_matches_eq(item.model.as_str(), model_filter.as_deref())
|
||||
admin_usage_matches_search(
|
||||
item,
|
||||
search.as_deref(),
|
||||
&users_by_id,
|
||||
&api_key_names,
|
||||
state.has_auth_user_data_reader(),
|
||||
state.has_auth_api_key_data_reader(),
|
||||
) && admin_usage_matches_username(
|
||||
item,
|
||||
username_filter.as_deref(),
|
||||
&users_by_id,
|
||||
state.has_auth_user_data_reader(),
|
||||
) && admin_usage_matches_eq(item.model.as_str(), model_filter.as_deref())
|
||||
&& admin_usage_matches_eq(
|
||||
item.provider_name.as_str(),
|
||||
provider_filter.as_deref(),
|
||||
@@ -153,24 +176,6 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
|
||||
});
|
||||
let total = usage.len();
|
||||
|
||||
let user_ids: Vec<String> = usage
|
||||
.iter()
|
||||
.filter_map(|item| item.user_id.clone())
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect();
|
||||
let users_by_id: BTreeMap<String, aether_data::repository::users::StoredUserSummary> =
|
||||
if state.has_user_data_reader() && !user_ids.is_empty() {
|
||||
state
|
||||
.list_users_by_ids(&user_ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|user| (user.id.clone(), user))
|
||||
.collect()
|
||||
} else {
|
||||
BTreeMap::new()
|
||||
};
|
||||
|
||||
let provider_key_names = admin_usage_provider_key_names(state, &usage).await?;
|
||||
|
||||
let records = usage
|
||||
@@ -182,6 +187,9 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
|
||||
return Ok(Some(build_admin_usage_records_response(
|
||||
&records,
|
||||
&users_by_id,
|
||||
&api_key_names,
|
||||
state.has_auth_user_data_reader(),
|
||||
state.has_auth_api_key_data_reader(),
|
||||
&provider_key_names,
|
||||
total,
|
||||
limit,
|
||||
|
||||
@@ -18,6 +18,23 @@ pub(crate) fn normalize_auth_type(value: Option<&str>) -> Result<String, String>
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_pool_advanced_config(
|
||||
value: Option<serde_json::Value>,
|
||||
) -> Result<Option<serde_json::Value>, String> {
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
match value {
|
||||
serde_json::Value::Null => Ok(None),
|
||||
// `pool_advanced: {}` still means "enable pool mode with defaults".
|
||||
serde_json::Value::Object(map) => Ok(Some(serde_json::Value::Object(map))),
|
||||
// Backward compatibility for older boolean payloads.
|
||||
serde_json::Value::Bool(true) => Ok(Some(serde_json::json!({}))),
|
||||
serde_json::Value::Bool(false) => Ok(None),
|
||||
_ => Err("pool_advanced 必须是 JSON 对象".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn validate_vertex_api_formats(
|
||||
provider_type: &str,
|
||||
auth_type: &str,
|
||||
@@ -46,3 +63,29 @@ pub(crate) fn validate_vertex_api_formats(
|
||||
allowed.join(", ")
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::normalize_pool_advanced_config;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn normalize_pool_advanced_preserves_empty_object() {
|
||||
assert_eq!(
|
||||
normalize_pool_advanced_config(Some(json!({}))).expect("empty object should normalize"),
|
||||
Some(json!({}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_pool_advanced_accepts_legacy_booleans() {
|
||||
assert_eq!(
|
||||
normalize_pool_advanced_config(Some(json!(true))).expect("true should normalize"),
|
||||
Some(json!({}))
|
||||
);
|
||||
assert_eq!(
|
||||
normalize_pool_advanced_config(Some(json!(false))).expect("false should normalize"),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,9 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderCreateReque
|
||||
use crate::handlers::admin::provider::shared::support::{
|
||||
normalize_provider_billing_type, parse_optional_rfc3339_unix_secs,
|
||||
};
|
||||
use crate::handlers::admin::provider::write::normalize::normalize_provider_type_input;
|
||||
use crate::handlers::admin::provider::write::normalize::{
|
||||
normalize_pool_advanced_config, normalize_provider_type_input,
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::handlers::admin::shared::normalize_json_object;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||
@@ -118,7 +120,7 @@ pub(crate) async fn build_admin_create_provider_record(
|
||||
let mut config_map = normalize_json_object(payload.config, "config")?
|
||||
.and_then(|value| value.as_object().cloned())
|
||||
.unwrap_or_default();
|
||||
if let Some(value) = normalize_json_object(payload.pool_advanced, "pool_advanced")? {
|
||||
if let Some(value) = normalize_pool_advanced_config(payload.pool_advanced)? {
|
||||
config_map.insert("pool_advanced".to_string(), value);
|
||||
}
|
||||
if let Some(value) = normalize_json_object(payload.failover_rules, "failover_rules")? {
|
||||
|
||||
@@ -2,7 +2,9 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderUpdatePatch
|
||||
use crate::handlers::admin::provider::shared::support::{
|
||||
normalize_provider_billing_type, parse_optional_rfc3339_unix_secs,
|
||||
};
|
||||
use crate::handlers::admin::provider::write::normalize::normalize_provider_type_input;
|
||||
use crate::handlers::admin::provider::write::normalize::{
|
||||
normalize_pool_advanced_config, normalize_provider_type_input,
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::handlers::admin::shared::normalize_json_object;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||
@@ -252,7 +254,7 @@ pub(crate) async fn build_admin_update_provider_record(
|
||||
if fields.is_null("pool_advanced") {
|
||||
config_map.remove("pool_advanced");
|
||||
} else {
|
||||
let value = normalize_json_object(payload.pool_advanced, "pool_advanced")?
|
||||
let value = normalize_pool_advanced_config(payload.pool_advanced)?
|
||||
.ok_or_else(|| "pool_advanced 必须是 JSON 对象".to_string())?;
|
||||
config_map.insert("pool_advanced".to_string(), value);
|
||||
}
|
||||
|
||||
@@ -84,6 +84,10 @@ impl<'a> AdminAppState<'a> {
|
||||
self.app.has_user_data_reader()
|
||||
}
|
||||
|
||||
pub(crate) fn has_auth_user_data_reader(&self) -> bool {
|
||||
self.app.has_auth_user_data_reader()
|
||||
}
|
||||
|
||||
pub(crate) fn has_proxy_node_reader(&self) -> bool {
|
||||
self.app.has_proxy_node_reader()
|
||||
}
|
||||
|
||||
@@ -54,6 +54,17 @@ impl<'a> AdminAppState<'a> {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_request_usage_body_ref(
|
||||
&self,
|
||||
body_ref: &str,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
self.app
|
||||
.data
|
||||
.resolve_request_usage_body_ref(body_ref)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_admin_usage_for_range(
|
||||
&self,
|
||||
time_range: &crate::handlers::admin::observability::AdminStatsTimeRange,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -16,6 +16,16 @@ impl<'a> AdminAppState<'a> {
|
||||
self.app.list_users_by_ids(user_ids).await
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_auth_user_summaries_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<
|
||||
std::collections::BTreeMap<String, aether_data::repository::users::StoredUserSummary>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.app.resolve_auth_user_summaries_by_ids(user_ids).await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_export_users_page(
|
||||
&self,
|
||||
query: &aether_data::repository::users::UserExportListQuery,
|
||||
@@ -103,6 +113,17 @@ impl<'a> AdminAppState<'a> {
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn initialize_auth_api_key_wallet(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
initial_gift_usd: f64,
|
||||
unlimited: bool,
|
||||
) -> Result<Option<aether_data::repository::wallet::StoredWalletSnapshot>, GatewayError> {
|
||||
self.app
|
||||
.initialize_auth_api_key_wallet(api_key_id, initial_gift_usd, unlimited)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_local_auth_user_profile(
|
||||
&self,
|
||||
user_id: &str,
|
||||
@@ -165,10 +186,94 @@ impl<'a> AdminAppState<'a> {
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_auth_api_key_wallet_limit_mode(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
limit_mode: &str,
|
||||
) -> Result<Option<aether_data::repository::wallet::StoredWalletSnapshot>, GatewayError> {
|
||||
self.app
|
||||
.update_auth_api_key_wallet_limit_mode(api_key_id, limit_mode)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn update_auth_user_wallet_snapshot(
|
||||
&self,
|
||||
user_id: &str,
|
||||
balance: f64,
|
||||
gift_balance: f64,
|
||||
limit_mode: &str,
|
||||
currency: &str,
|
||||
status: &str,
|
||||
total_recharged: f64,
|
||||
total_consumed: f64,
|
||||
total_refunded: f64,
|
||||
total_adjusted: f64,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Result<Option<aether_data::repository::wallet::StoredWalletSnapshot>, GatewayError> {
|
||||
self.app
|
||||
.update_auth_user_wallet_snapshot(
|
||||
user_id,
|
||||
balance,
|
||||
gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
total_recharged,
|
||||
total_consumed,
|
||||
total_refunded,
|
||||
total_adjusted,
|
||||
updated_at_unix_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn update_auth_api_key_wallet_snapshot(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
balance: f64,
|
||||
gift_balance: f64,
|
||||
limit_mode: &str,
|
||||
currency: &str,
|
||||
status: &str,
|
||||
total_recharged: f64,
|
||||
total_consumed: f64,
|
||||
total_refunded: f64,
|
||||
total_adjusted: f64,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Result<Option<aether_data::repository::wallet::StoredWalletSnapshot>, GatewayError> {
|
||||
self.app
|
||||
.update_auth_api_key_wallet_snapshot(
|
||||
api_key_id,
|
||||
balance,
|
||||
gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
total_recharged,
|
||||
total_consumed,
|
||||
total_refunded,
|
||||
total_adjusted,
|
||||
updated_at_unix_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn count_active_admin_users(&self) -> Result<u64, GatewayError> {
|
||||
self.app.count_active_admin_users().await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_user_model_capability_settings(
|
||||
&self,
|
||||
user_id: &str,
|
||||
settings: Option<serde_json::Value>,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
self.app
|
||||
.update_user_model_capability_settings(user_id, settings)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn count_user_pending_refunds(
|
||||
&self,
|
||||
user_id: &str,
|
||||
@@ -236,6 +341,24 @@ impl<'a> AdminAppState<'a> {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_auth_api_key_snapshots_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::auth::StoredAuthApiKeySnapshot>, GatewayError> {
|
||||
self.app
|
||||
.resolve_auth_api_key_snapshots_by_ids(api_key_ids)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_auth_api_key_names_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<std::collections::BTreeMap<String, String>, GatewayError> {
|
||||
self.app
|
||||
.resolve_auth_api_key_names_by_ids(api_key_ids)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_auth_api_key_export_records_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
@@ -339,6 +462,18 @@ impl<'a> AdminAppState<'a> {
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn set_user_api_key_active(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
is_active: bool,
|
||||
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
|
||||
{
|
||||
self.app
|
||||
.set_user_api_key_active(user_id, api_key_id, is_active)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn set_user_api_key_locked(
|
||||
&self,
|
||||
user_id: &str,
|
||||
@@ -362,6 +497,18 @@ impl<'a> AdminAppState<'a> {
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn set_user_api_key_force_capabilities(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
force_capabilities: Option<serde_json::Value>,
|
||||
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
|
||||
{
|
||||
self.app
|
||||
.set_user_api_key_force_capabilities(user_id, api_key_id, force_capabilities)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_user_api_key(
|
||||
&self,
|
||||
user_id: &str,
|
||||
|
||||
@@ -137,11 +137,46 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
|
||||
)));
|
||||
}
|
||||
|
||||
if decision.route_kind.as_deref() == Some("users_import")
|
||||
&& request_method == http::Method::POST
|
||||
&& request_path == "/api/admin/system/users/import"
|
||||
{
|
||||
let Some(request_body) = request_body else {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": "请求数据验证失败" })),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
};
|
||||
return Ok(Some(
|
||||
match state
|
||||
.import_admin_system_users(
|
||||
request_body,
|
||||
decision
|
||||
.admin_principal
|
||||
.as_ref()
|
||||
.map(|principal| principal.user_id.as_str()),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Ok(payload) => attach_admin_audit_response(
|
||||
Json(payload).into_response(),
|
||||
"admin_system_users_imported",
|
||||
"import_system_users",
|
||||
"system_users_import",
|
||||
"global",
|
||||
),
|
||||
Err((status, payload)) => (status, Json(payload)).into_response(),
|
||||
},
|
||||
));
|
||||
}
|
||||
|
||||
if matches!(
|
||||
decision.route_kind.as_deref(),
|
||||
Some(
|
||||
"users_import"
|
||||
| "smtp_test"
|
||||
"smtp_test"
|
||||
| "cleanup"
|
||||
| "purge_config"
|
||||
| "purge_users"
|
||||
|
||||
@@ -120,8 +120,17 @@ pub(crate) async fn build_admin_create_user_api_key_response(
|
||||
key_hash: hash_admin_user_api_key(&plaintext_key),
|
||||
key_encrypted: Some(key_encrypted),
|
||||
name: Some(name.clone()),
|
||||
allowed_providers: None,
|
||||
allowed_api_formats: None,
|
||||
allowed_models: None,
|
||||
rate_limit,
|
||||
concurrent_limit: 5,
|
||||
force_capabilities: None,
|
||||
is_active: true,
|
||||
expires_at_unix_secs: None,
|
||||
auto_delete_on_expiry: false,
|
||||
total_requests: 0,
|
||||
total_cost_usd: 0.0,
|
||||
})
|
||||
.await?
|
||||
else {
|
||||
|
||||
@@ -31,8 +31,9 @@ use crate::control::{
|
||||
GatewayPublicRequestContext,
|
||||
};
|
||||
use crate::executor::{
|
||||
maybe_execute_stream_request, maybe_execute_sync_request,
|
||||
record_failed_usage_for_exhausted_request, LocalExecutionRequestOutcome,
|
||||
build_local_execution_runtime_miss_context, maybe_execute_stream_request,
|
||||
maybe_execute_sync_request, record_failed_usage_for_exhausted_request,
|
||||
record_failed_usage_for_runtime_miss_request, LocalExecutionRequestOutcome,
|
||||
};
|
||||
use crate::frontdoor_loop_guard::{
|
||||
frontdoor_self_loop_public_ai_path, request_has_execution_runtime_loop_guard,
|
||||
@@ -46,7 +47,7 @@ use crate::headers::{extract_or_generate_trace_id, should_skip_request_header};
|
||||
use crate::router::RequestAdmissionError;
|
||||
use crate::{
|
||||
AppState, FrontdoorUserRpmOutcome, GatewayError, GatewayFallbackMetricKind,
|
||||
GatewayFallbackReason,
|
||||
GatewayFallbackReason, LocalExecutionRuntimeMissDiagnostic,
|
||||
};
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::extract::{ConnectInfo, Request, State};
|
||||
@@ -1013,9 +1014,16 @@ pub(crate) async fn proxy_request(
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
}
|
||||
let local_execution_runtime_miss_detail =
|
||||
local_execution_runtime_miss_detail(control_decision)
|
||||
.unwrap_or("AI public execution runtime miss did not match a Rust execution path");
|
||||
let local_execution_runtime_miss_diagnostic =
|
||||
state.take_local_execution_runtime_miss_diagnostic(&trace_id);
|
||||
let local_execution_runtime_miss_detail = local_execution_runtime_miss_detail(
|
||||
control_decision,
|
||||
local_execution_runtime_miss_diagnostic.as_ref(),
|
||||
stream_request,
|
||||
)
|
||||
.unwrap_or_else(|| {
|
||||
"AI public execution runtime miss did not match a Rust execution path".to_string()
|
||||
});
|
||||
state.record_fallback_metric(
|
||||
GatewayFallbackMetricKind::LocalExecutionRuntimeMiss,
|
||||
control_decision,
|
||||
@@ -1023,38 +1031,95 @@ pub(crate) async fn proxy_request(
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS),
|
||||
GatewayFallbackReason::LocalExecutionPathRequired,
|
||||
);
|
||||
let local_execution_runtime_miss_diagnostic =
|
||||
state.take_local_execution_runtime_miss_diagnostic(&trace_id);
|
||||
if let Some(diagnostic) = local_execution_runtime_miss_diagnostic.as_ref() {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
local_execution_runtime_miss_reason = %diagnostic.reason,
|
||||
route_family = diagnostic.route_family.as_deref().unwrap_or_default(),
|
||||
route_kind = diagnostic.route_kind.as_deref().unwrap_or_default(),
|
||||
public_path = diagnostic.public_path.as_deref().unwrap_or_default(),
|
||||
plan_kind = diagnostic.plan_kind.as_deref().unwrap_or_default(),
|
||||
requested_model = diagnostic.requested_model.as_deref().unwrap_or_default(),
|
||||
candidate_count = diagnostic.candidate_count.unwrap_or(0),
|
||||
skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0),
|
||||
skip_reasons = diagnostic.skip_reasons_summary().unwrap_or_default(),
|
||||
"gateway local execution runtime miss"
|
||||
);
|
||||
}
|
||||
let local_execution_runtime_miss_context =
|
||||
build_local_execution_runtime_miss_context(&state, &trace_id, control_decision).await;
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
local_execution_runtime_miss_reason = local_execution_runtime_miss_diagnostic
|
||||
.as_ref()
|
||||
.map(|value| value.reason.as_str())
|
||||
.unwrap_or("unknown"),
|
||||
route_family = local_execution_runtime_miss_diagnostic
|
||||
.as_ref()
|
||||
.and_then(|value| value.route_family.as_deref())
|
||||
.or_else(|| control_decision.and_then(|value| value.route_family.as_deref()))
|
||||
.unwrap_or_default(),
|
||||
route_kind = local_execution_runtime_miss_diagnostic
|
||||
.as_ref()
|
||||
.and_then(|value| value.route_kind.as_deref())
|
||||
.or_else(|| control_decision.and_then(|value| value.route_kind.as_deref()))
|
||||
.unwrap_or_default(),
|
||||
public_path = local_execution_runtime_miss_diagnostic
|
||||
.as_ref()
|
||||
.and_then(|value| value.public_path.as_deref())
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| control_decision.map(GatewayControlDecision::proxy_path_and_query))
|
||||
.unwrap_or_default(),
|
||||
plan_kind = local_execution_runtime_miss_diagnostic
|
||||
.as_ref()
|
||||
.and_then(|value| value.plan_kind.as_deref())
|
||||
.unwrap_or_default(),
|
||||
requested_model = local_execution_runtime_miss_diagnostic
|
||||
.as_ref()
|
||||
.and_then(|value| value.requested_model.as_deref())
|
||||
.unwrap_or_default(),
|
||||
candidate_count = local_execution_runtime_miss_diagnostic
|
||||
.as_ref()
|
||||
.and_then(|value| value.candidate_count)
|
||||
.unwrap_or(0),
|
||||
persisted_candidate_count = local_execution_runtime_miss_context.persisted_candidate_count(),
|
||||
skipped_candidate_count = local_execution_runtime_miss_diagnostic
|
||||
.as_ref()
|
||||
.and_then(|value| value.skipped_candidate_count)
|
||||
.unwrap_or(0),
|
||||
skip_reasons = local_execution_runtime_miss_diagnostic
|
||||
.as_ref()
|
||||
.and_then(|value| value.skip_reasons_summary())
|
||||
.unwrap_or_default(),
|
||||
auth_user_id = local_execution_runtime_miss_context
|
||||
.auth_user_id
|
||||
.as_deref()
|
||||
.unwrap_or_default(),
|
||||
auth_api_key_id = local_execution_runtime_miss_context
|
||||
.auth_api_key_id
|
||||
.as_deref()
|
||||
.unwrap_or_default(),
|
||||
auth_api_key_name = local_execution_runtime_miss_context
|
||||
.auth_api_key_name
|
||||
.as_deref()
|
||||
.unwrap_or_default(),
|
||||
request_candidates = local_execution_runtime_miss_context
|
||||
.candidate_summary()
|
||||
.unwrap_or_default(),
|
||||
"gateway local execution runtime miss"
|
||||
);
|
||||
if let Some(exhaustion) = local_execution_exhaustion {
|
||||
record_failed_usage_for_exhausted_request(
|
||||
&state,
|
||||
exhaustion,
|
||||
&started_at,
|
||||
local_execution_runtime_miss_detail,
|
||||
local_execution_runtime_miss_detail.as_str(),
|
||||
local_execution_runtime_miss_diagnostic.as_ref(),
|
||||
)
|
||||
.await;
|
||||
} else {
|
||||
record_failed_usage_for_runtime_miss_request(
|
||||
&state,
|
||||
&trace_id,
|
||||
&started_at,
|
||||
local_execution_runtime_miss_detail.as_str(),
|
||||
control_decision,
|
||||
local_execution_runtime_miss_diagnostic.as_ref(),
|
||||
&local_execution_runtime_miss_context,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
let mut response = build_local_http_error_response(
|
||||
&trace_id,
|
||||
control_decision,
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
local_execution_runtime_miss_detail,
|
||||
local_execution_runtime_miss_detail.as_str(),
|
||||
)?;
|
||||
if let Some(diagnostic) = local_execution_runtime_miss_diagnostic {
|
||||
if !diagnostic.reason.trim().is_empty() {
|
||||
@@ -1095,6 +1160,41 @@ pub(crate) async fn proxy_request(
|
||||
|
||||
fn local_execution_runtime_miss_detail(
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
diagnostic: Option<&LocalExecutionRuntimeMissDiagnostic>,
|
||||
stream_request: bool,
|
||||
) -> Option<String> {
|
||||
if let Some(detail) = local_execution_runtime_miss_model_detail(diagnostic, stream_request) {
|
||||
return Some(detail);
|
||||
}
|
||||
|
||||
local_execution_runtime_miss_route_detail(decision).map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn local_execution_runtime_miss_model_detail(
|
||||
diagnostic: Option<&LocalExecutionRuntimeMissDiagnostic>,
|
||||
stream_request: bool,
|
||||
) -> Option<String> {
|
||||
let diagnostic = diagnostic?;
|
||||
if !matches!(
|
||||
diagnostic.reason.as_str(),
|
||||
"candidate_list_empty" | "all_candidates_skipped"
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
|
||||
let requested_model = diagnostic
|
||||
.requested_model
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
let request_mode = if stream_request { "流式" } else { "同步" };
|
||||
Some(format!(
|
||||
"没有可用的提供商支持模型 {requested_model} 的{request_mode}请求"
|
||||
))
|
||||
}
|
||||
|
||||
fn local_execution_runtime_miss_route_detail(
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
) -> Option<&'static str> {
|
||||
let decision = decision?;
|
||||
if decision.route_class.as_deref() != Some("ai_public") {
|
||||
@@ -1121,6 +1221,60 @@ fn local_execution_runtime_miss_detail(
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
local_execution_runtime_miss_detail, GatewayControlDecision,
|
||||
LocalExecutionRuntimeMissDiagnostic,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn runtime_miss_detail_returns_model_specific_stream_message_when_candidates_are_unavailable() {
|
||||
let decision = GatewayControlDecision::synthetic(
|
||||
"/v1/chat/completions",
|
||||
Some("ai_public".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
);
|
||||
let diagnostic = LocalExecutionRuntimeMissDiagnostic {
|
||||
reason: "candidate_list_empty".to_string(),
|
||||
requested_model: Some("gpt-5.4".to_string()),
|
||||
..LocalExecutionRuntimeMissDiagnostic::default()
|
||||
};
|
||||
|
||||
let detail = local_execution_runtime_miss_detail(Some(&decision), Some(&diagnostic), true);
|
||||
|
||||
assert_eq!(
|
||||
detail.as_deref(),
|
||||
Some("没有可用的提供商支持模型 gpt-5.4 的流式请求")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_miss_detail_falls_back_to_route_default_when_reason_is_not_model_unavailable() {
|
||||
let decision = GatewayControlDecision::synthetic(
|
||||
"/v1/messages",
|
||||
Some("ai_public".to_string()),
|
||||
Some("claude".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("claude:chat".to_string()),
|
||||
);
|
||||
let diagnostic = LocalExecutionRuntimeMissDiagnostic {
|
||||
reason: "missing_auth_context".to_string(),
|
||||
requested_model: Some("claude-sonnet-4-5".to_string()),
|
||||
..LocalExecutionRuntimeMissDiagnostic::default()
|
||||
};
|
||||
|
||||
let detail = local_execution_runtime_miss_detail(Some(&decision), Some(&diagnostic), false);
|
||||
|
||||
assert_eq!(
|
||||
detail.as_deref(),
|
||||
Some("Claude messages execution runtime miss did not match a Rust execution path")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[path = "finalize.rs"]
|
||||
mod finalize;
|
||||
|
||||
|
||||
@@ -1060,23 +1060,21 @@ pub(super) async fn handle_dashboard_recent_requests_get(
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
let users_by_id: BTreeMap<String, aether_data::repository::users::StoredUserSummary> =
|
||||
if state.has_user_data_reader() && !user_ids.is_empty() {
|
||||
match state.list_users_by_ids(&user_ids).await {
|
||||
Ok(value) => value
|
||||
.into_iter()
|
||||
.map(|user| (user.id.clone(), user))
|
||||
.collect(),
|
||||
Err(err) => {
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("dashboard user lookup failed: {err:?}"),
|
||||
false,
|
||||
)
|
||||
}
|
||||
match state.resolve_auth_user_summaries_by_ids(&user_ids).await {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("dashboard user lookup failed: {err:?}"),
|
||||
false,
|
||||
)
|
||||
}
|
||||
} else {
|
||||
BTreeMap::new()
|
||||
};
|
||||
let mut usernames_by_id: BTreeMap<String, String> = users_by_id
|
||||
.iter()
|
||||
.filter(|(_, user)| !user.username.trim().is_empty())
|
||||
.map(|(user_id, user)| (user_id.clone(), user.username.clone()))
|
||||
.collect();
|
||||
|
||||
let requests = usage
|
||||
.into_iter()
|
||||
@@ -1084,9 +1082,13 @@ pub(super) async fn handle_dashboard_recent_requests_get(
|
||||
let username = item
|
||||
.user_id
|
||||
.as_ref()
|
||||
.and_then(|user_id| users_by_id.get(user_id))
|
||||
.map(|user| user.username.clone())
|
||||
.or(item.username.clone())
|
||||
.and_then(|user_id| usernames_by_id.get(user_id))
|
||||
.cloned()
|
||||
.or_else(|| {
|
||||
(!state.has_auth_user_data_reader())
|
||||
.then(|| item.username.clone())
|
||||
.flatten()
|
||||
})
|
||||
.unwrap_or_else(|| "Unknown".to_string());
|
||||
json!({
|
||||
"id": item.id,
|
||||
|
||||
@@ -530,8 +530,17 @@ pub(super) async fn handle_users_me_api_key_create(
|
||||
key_hash: hash_users_me_api_key(&plaintext_key),
|
||||
key_encrypted: Some(key_encrypted),
|
||||
name: Some(name.clone()),
|
||||
allowed_providers: None,
|
||||
allowed_api_formats: None,
|
||||
allowed_models: None,
|
||||
rate_limit,
|
||||
concurrent_limit: 5,
|
||||
force_capabilities: None,
|
||||
is_active: true,
|
||||
expires_at_unix_secs: None,
|
||||
auto_delete_on_expiry: false,
|
||||
total_requests: 0,
|
||||
total_cost_usd: 0.0,
|
||||
};
|
||||
let Some(created) = (match state.create_user_api_key(record).await {
|
||||
Ok(value) => value,
|
||||
|
||||
@@ -145,32 +145,56 @@ fn users_me_usage_cache_hit_rate(total_input_context: u64, cache_read_tokens: u6
|
||||
}
|
||||
}
|
||||
|
||||
fn users_me_usage_matches_search(item: &StoredRequestUsageAudit, search: Option<&str>) -> bool {
|
||||
fn users_me_usage_api_key_name(
|
||||
item: &StoredRequestUsageAudit,
|
||||
api_key_names: &BTreeMap<String, String>,
|
||||
auth_api_key_reader_available: bool,
|
||||
) -> Option<String> {
|
||||
item.api_key_id
|
||||
.as_ref()
|
||||
.and_then(|api_key_id| api_key_names.get(api_key_id))
|
||||
.cloned()
|
||||
.or_else(|| {
|
||||
(!auth_api_key_reader_available)
|
||||
.then(|| item.api_key_name.clone())
|
||||
.flatten()
|
||||
})
|
||||
}
|
||||
|
||||
fn users_me_usage_matches_search(
|
||||
item: &StoredRequestUsageAudit,
|
||||
search: Option<&str>,
|
||||
api_key_names: &BTreeMap<String, String>,
|
||||
auth_api_key_reader_available: bool,
|
||||
) -> bool {
|
||||
let Some(search) = search.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||
return true;
|
||||
};
|
||||
|
||||
let model = item.model.to_ascii_lowercase();
|
||||
let api_key_name = item
|
||||
.api_key_name
|
||||
.as_deref()
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase();
|
||||
let api_key_name =
|
||||
users_me_usage_api_key_name(item, api_key_names, auth_api_key_reader_available)
|
||||
.as_deref()
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase();
|
||||
search.split_whitespace().all(|keyword| {
|
||||
let keyword = keyword.to_ascii_lowercase();
|
||||
model.contains(&keyword) || api_key_name.contains(&keyword)
|
||||
})
|
||||
}
|
||||
|
||||
fn build_users_me_usage_api_key_payload(item: &StoredRequestUsageAudit) -> serde_json::Value {
|
||||
fn build_users_me_usage_api_key_payload(
|
||||
item: &StoredRequestUsageAudit,
|
||||
api_key_names: &BTreeMap<String, String>,
|
||||
auth_api_key_reader_available: bool,
|
||||
) -> serde_json::Value {
|
||||
let api_key_name =
|
||||
users_me_usage_api_key_name(item, api_key_names, auth_api_key_reader_available);
|
||||
match item.api_key_id.as_deref() {
|
||||
Some(api_key_id) => json!({
|
||||
"id": api_key_id,
|
||||
"name": item.api_key_name.clone(),
|
||||
"display": item
|
||||
.api_key_name
|
||||
.clone()
|
||||
.unwrap_or_else(|| api_key_id.to_string()),
|
||||
"name": api_key_name.clone(),
|
||||
"display": api_key_name.unwrap_or_else(|| api_key_id.to_string()),
|
||||
}),
|
||||
None => serde_json::Value::Null,
|
||||
}
|
||||
@@ -179,7 +203,14 @@ fn build_users_me_usage_api_key_payload(item: &StoredRequestUsageAudit) -> serde
|
||||
fn build_users_me_usage_record_payload(
|
||||
item: &StoredRequestUsageAudit,
|
||||
include_actual_cost: bool,
|
||||
api_key_names: &BTreeMap<String, String>,
|
||||
auth_api_key_reader_available: bool,
|
||||
) -> serde_json::Value {
|
||||
let input_price_per_1m = item.settlement_input_price_per_1m();
|
||||
let output_price_per_1m = item.settlement_output_price_per_1m();
|
||||
let cache_creation_price_per_1m = item.settlement_cache_creation_price_per_1m();
|
||||
let cache_read_price_per_1m = item.settlement_cache_read_price_per_1m();
|
||||
let rate_multiplier = item.settlement_rate_multiplier();
|
||||
let mut payload = json!({
|
||||
"id": item.id,
|
||||
"model": item.model,
|
||||
@@ -203,11 +234,15 @@ fn build_users_me_usage_record_payload(
|
||||
"cache_read_input_tokens": item.cache_read_input_tokens,
|
||||
"status_code": item.status_code,
|
||||
"error_message": item.error_message,
|
||||
"input_price_per_1m": serde_json::Value::Null,
|
||||
"output_price_per_1m": item.output_price_per_1m,
|
||||
"cache_creation_price_per_1m": serde_json::Value::Null,
|
||||
"cache_read_price_per_1m": serde_json::Value::Null,
|
||||
"api_key": build_users_me_usage_api_key_payload(item),
|
||||
"input_price_per_1m": input_price_per_1m,
|
||||
"output_price_per_1m": output_price_per_1m,
|
||||
"cache_creation_price_per_1m": cache_creation_price_per_1m,
|
||||
"cache_read_price_per_1m": cache_read_price_per_1m,
|
||||
"api_key": build_users_me_usage_api_key_payload(
|
||||
item,
|
||||
api_key_names,
|
||||
auth_api_key_reader_available,
|
||||
),
|
||||
});
|
||||
|
||||
if item.target_model.is_some() {
|
||||
@@ -215,7 +250,7 @@ fn build_users_me_usage_record_payload(
|
||||
}
|
||||
if include_actual_cost {
|
||||
payload["actual_cost"] = json!(round_to(item.actual_total_cost_usd, 6));
|
||||
payload["rate_multiplier"] = serde_json::Value::Null;
|
||||
payload["rate_multiplier"] = json!(rate_multiplier);
|
||||
}
|
||||
payload
|
||||
}
|
||||
@@ -233,7 +268,7 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
|
||||
"cache_read_input_tokens": item.cache_read_input_tokens,
|
||||
"cost": round_to(item.total_cost_usd, 6),
|
||||
"actual_cost": round_to(item.actual_total_cost_usd, 6),
|
||||
"rate_multiplier": serde_json::Value::Null,
|
||||
"rate_multiplier": item.settlement_rate_multiplier(),
|
||||
"response_time_ms": item.response_time_ms,
|
||||
"first_byte_time_ms": item.first_byte_time_ms,
|
||||
"api_format": item.api_format,
|
||||
@@ -675,6 +710,30 @@ pub(super) async fn handle_users_me_usage_get(
|
||||
)
|
||||
}
|
||||
};
|
||||
let api_key_names = if state.has_auth_api_key_data_reader() {
|
||||
let api_key_ids = usage
|
||||
.iter()
|
||||
.filter_map(|item| item.api_key_id.clone())
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
if api_key_ids.is_empty() {
|
||||
BTreeMap::new()
|
||||
} else {
|
||||
match state.resolve_auth_api_key_names_by_ids(&api_key_ids).await {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("user api key name lookup failed: {err:?}"),
|
||||
false,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
BTreeMap::new()
|
||||
};
|
||||
|
||||
let summary_items = usage
|
||||
.iter()
|
||||
@@ -738,7 +797,14 @@ pub(super) async fn handle_users_me_usage_get(
|
||||
|
||||
let mut records = usage
|
||||
.into_iter()
|
||||
.filter(|item| users_me_usage_matches_search(item, search.as_deref()))
|
||||
.filter(|item| {
|
||||
users_me_usage_matches_search(
|
||||
item,
|
||||
search.as_deref(),
|
||||
&api_key_names,
|
||||
state.has_auth_api_key_data_reader(),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
records.sort_by(|left, right| {
|
||||
right
|
||||
@@ -751,7 +817,14 @@ pub(super) async fn handle_users_me_usage_get(
|
||||
.into_iter()
|
||||
.skip(offset)
|
||||
.take(limit)
|
||||
.map(|item| build_users_me_usage_record_payload(&item, include_actual_cost))
|
||||
.map(|item| {
|
||||
build_users_me_usage_record_payload(
|
||||
&item,
|
||||
include_actual_cost,
|
||||
&api_key_names,
|
||||
state.has_auth_api_key_data_reader(),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let wallet = state
|
||||
|
||||
@@ -245,6 +245,7 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
|
||||
| (Some("provider_oauth_manage"), http::Method::POST, Some("device_authorize"))
|
||||
| (Some("provider_oauth_manage"), http::Method::POST, Some("device_poll"))
|
||||
| (Some("system_manage"), http::Method::POST, Some("config_import"))
|
||||
| (Some("system_manage"), http::Method::POST, Some("users_import"))
|
||||
| (Some("system_manage"), http::Method::PUT, Some("settings_set"))
|
||||
| (Some("system_manage"), http::Method::PUT, Some("config_set"))
|
||||
| (Some("system_manage"), http::Method::PUT, Some("email_template_set"))
|
||||
|
||||
@@ -137,6 +137,8 @@ impl From<GatewayLogRotationArg> for LogRotation {
|
||||
}
|
||||
}
|
||||
|
||||
const GATEWAY_TOKIO_WORKER_STACK_SIZE_BYTES: usize = 8 * 1024 * 1024;
|
||||
|
||||
fn env_var_trimmed(name: &str) -> Option<String> {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
@@ -527,6 +529,9 @@ struct Args {
|
||||
)]
|
||||
node_role: NodeRoleArg,
|
||||
|
||||
#[arg(long, default_value_t = false)]
|
||||
migrate: bool,
|
||||
|
||||
/// Path to frontend static files directory (SPA). When set, the gateway
|
||||
/// serves the frontend directly without nginx.
|
||||
#[arg(long, env = "AETHER_GATEWAY_STATIC_DIR")]
|
||||
@@ -608,11 +613,16 @@ struct Args {
|
||||
|
||||
impl Args {
|
||||
fn runtime_config(&self) -> Result<ServiceRuntimeConfig, std::io::Error> {
|
||||
let default_log_filter = if self.migrate {
|
||||
"aether_gateway=info,aether_data=info"
|
||||
} else {
|
||||
"aether_gateway=info"
|
||||
};
|
||||
let config = self
|
||||
.logging
|
||||
.apply_to_runtime_config(ServiceRuntimeConfig::new(
|
||||
"aether-gateway",
|
||||
"aether_gateway=info",
|
||||
default_log_filter,
|
||||
))?;
|
||||
Ok(config
|
||||
.with_node_role(self.node_role.as_str())
|
||||
@@ -741,9 +751,20 @@ fn validate_deployment_topology(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.thread_stack_size(GATEWAY_TOKIO_WORKER_STACK_SIZE_BYTES)
|
||||
.build()?
|
||||
.block_on(run())
|
||||
}
|
||||
|
||||
async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let args = Args::parse();
|
||||
if args.migrate {
|
||||
init_service_runtime(args.runtime_config()?)?;
|
||||
return run_explicit_migrations(&args).await;
|
||||
}
|
||||
let app_port = validate_app_port(args.app_port)?;
|
||||
let bind_addr = gateway_bind_addr(app_port)?;
|
||||
set_gateway_frontdoor_app_port(app_port);
|
||||
@@ -899,11 +920,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
execution_runtime_configured = state.execution_runtime_configured(),
|
||||
"aether-gateway data layer configured"
|
||||
);
|
||||
// Run pending database migrations before serving traffic
|
||||
info!("running database migrations...");
|
||||
if state.run_postgres_migrations().await? {
|
||||
info!("database migrations complete");
|
||||
}
|
||||
ensure_postgres_schema_is_current(&state).await?;
|
||||
let reset_stale_proxy_nodes = state.reset_stale_proxy_node_tunnel_statuses().await?;
|
||||
if reset_stale_proxy_nodes > 0 {
|
||||
info!(
|
||||
@@ -965,9 +982,161 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_explicit_migrations(args: &Args) -> Result<(), Box<dyn std::error::Error>> {
|
||||
if args.data.effective_postgres_url().is_none() {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
"AETHER_GATEWAY_DATA_POSTGRES_URL or DATABASE_URL is required when running --migrate",
|
||||
)
|
||||
.into());
|
||||
}
|
||||
|
||||
if args.data.configured_encryption_key_mismatch() {
|
||||
warn!(
|
||||
"AETHER_GATEWAY_DATA_ENCRYPTION_KEY differs from ENCRYPTION_KEY; aether-gateway will prefer the gateway-specific value"
|
||||
);
|
||||
}
|
||||
|
||||
let state = AppState::new()?.with_data_config(args.data.to_config())?;
|
||||
let pending = state
|
||||
.pending_postgres_migrations()
|
||||
.await?
|
||||
.unwrap_or_default();
|
||||
if pending.is_empty() {
|
||||
info!(
|
||||
pending_migrations = 0,
|
||||
"database migrations already up to date"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let next = pending
|
||||
.first()
|
||||
.expect("pending migrations should have a first element");
|
||||
info!(
|
||||
pending_migrations = pending.len(),
|
||||
next_version = next.version,
|
||||
next_description = %next.description,
|
||||
pending_versions = %format_pending_migrations(&pending),
|
||||
"running database migrations by explicit request..."
|
||||
);
|
||||
if state.run_postgres_migrations().await? {
|
||||
info!("database migrations complete");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn format_pending_migrations(pending: &[aether_data::migrate::PendingMigrationInfo]) -> String {
|
||||
pending
|
||||
.iter()
|
||||
.map(|migration| format!("{} ({})", migration.version, migration.description))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
}
|
||||
|
||||
async fn ensure_postgres_schema_is_current(
|
||||
state: &AppState,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let Some(pending) = state.prepare_postgres_for_startup().await? else {
|
||||
return Ok(());
|
||||
};
|
||||
if pending.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let next = pending
|
||||
.first()
|
||||
.expect("pending migrations should have a first element");
|
||||
Err(pending_schema_error(pending.len(), next.version, &next.description).into())
|
||||
}
|
||||
|
||||
fn pending_schema_error(
|
||||
pending_count: usize,
|
||||
next_version: i64,
|
||||
next_description: &str,
|
||||
) -> std::io::Error {
|
||||
std::io::Error::other(format!(
|
||||
"database schema is behind by {} migration(s); next pending migration is {} ({})\nrun `aether-gateway --migrate` before starting the service",
|
||||
pending_count,
|
||||
next_version,
|
||||
next_description
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::resolve_healthcheck_url;
|
||||
use super::{
|
||||
ensure_postgres_schema_is_current, pending_schema_error, resolve_healthcheck_url, Args,
|
||||
DeploymentTopologyArg, GatewayDataArgs, GatewayFrontdoorArgs, GatewayLogDestinationArg,
|
||||
GatewayLogFormatArg, GatewayLogRotationArg, GatewayLoggingArgs, GatewayRateLimitArgs,
|
||||
GatewayUsageArgs, NodeRoleArg, VideoTaskTruthSourceArg,
|
||||
};
|
||||
use aether_gateway::AppState;
|
||||
|
||||
fn test_args() -> Args {
|
||||
Args {
|
||||
app_port: 8084,
|
||||
healthcheck: false,
|
||||
healthcheck_timeout_ms: 3_000,
|
||||
deployment_topology: DeploymentTopologyArg::SingleNode,
|
||||
node_role: NodeRoleArg::All,
|
||||
migrate: false,
|
||||
static_dir: None,
|
||||
video_task_truth_source_mode: VideoTaskTruthSourceArg::PythonSyncReport,
|
||||
video_task_poller_interval_ms: 5_000,
|
||||
video_task_poller_batch_size: 32,
|
||||
video_task_store_path: None,
|
||||
max_in_flight_requests: None,
|
||||
distributed_request_limit: None,
|
||||
distributed_request_redis_url: None,
|
||||
distributed_request_redis_key_prefix: None,
|
||||
distributed_request_lease_ttl_ms: 30_000,
|
||||
distributed_request_renew_interval_ms: 10_000,
|
||||
distributed_request_command_timeout_ms: 1_000,
|
||||
data: GatewayDataArgs {
|
||||
postgres_url: None,
|
||||
encryption_key: None,
|
||||
redis_url: None,
|
||||
redis_key_prefix: None,
|
||||
postgres_min_connections: 1,
|
||||
postgres_max_connections: 30,
|
||||
postgres_acquire_timeout_ms: 3_000,
|
||||
postgres_idle_timeout_ms: 60_000,
|
||||
postgres_max_lifetime_ms: 1_800_000,
|
||||
postgres_statement_cache_capacity: 100,
|
||||
postgres_require_ssl: false,
|
||||
},
|
||||
usage: GatewayUsageArgs {
|
||||
queue_stream_key: "usage:events".to_string(),
|
||||
queue_group: "usage_consumers".to_string(),
|
||||
queue_dlq_stream_key: "usage:events:dlq".to_string(),
|
||||
queue_stream_maxlen: 2_000,
|
||||
queue_batch_size: 200,
|
||||
queue_block_ms: 500,
|
||||
queue_reclaim_idle_ms: 30_000,
|
||||
queue_reclaim_count: 200,
|
||||
queue_reclaim_interval_ms: 5_000,
|
||||
},
|
||||
frontdoor: GatewayFrontdoorArgs {
|
||||
environment: "development".to_string(),
|
||||
cors_origins: None,
|
||||
cors_allow_credentials: true,
|
||||
},
|
||||
rate_limit: GatewayRateLimitArgs {
|
||||
bucket_seconds: 60,
|
||||
key_ttl_seconds: 120,
|
||||
fail_open: true,
|
||||
},
|
||||
logging: GatewayLoggingArgs {
|
||||
log_format: GatewayLogFormatArg::Pretty,
|
||||
log_destination: GatewayLogDestinationArg::Stdout,
|
||||
log_dir: None,
|
||||
log_rotation: GatewayLogRotationArg::Daily,
|
||||
log_retention_days: 7,
|
||||
log_max_files: 30,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_healthcheck_url_from_app_port() {
|
||||
@@ -982,4 +1151,65 @@ mod tests {
|
||||
let error = resolve_healthcheck_url(0).unwrap_err();
|
||||
assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_migrate_runtime_config_enables_data_logs() {
|
||||
let mut args = test_args();
|
||||
args.migrate = true;
|
||||
let config = args.runtime_config().expect("runtime config should build");
|
||||
assert_eq!(
|
||||
config.default_log_filter,
|
||||
"aether_gateway=info,aether_data=info"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normal_runtime_config_keeps_gateway_only_logs() {
|
||||
let config = test_args()
|
||||
.runtime_config()
|
||||
.expect("runtime config should build");
|
||||
assert_eq!(config.default_log_filter, "aether_gateway=info");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_schema_error_mentions_explicit_migrate_command() {
|
||||
let error = pending_schema_error(2, 20260413020000, "squash usage schema split");
|
||||
let message = error.to_string();
|
||||
assert!(message.contains("database schema is behind by 2 migration(s)"));
|
||||
assert!(message.contains("20260413020000"));
|
||||
assert!(message.contains("squash usage schema split"));
|
||||
assert!(message.contains("aether-gateway --migrate"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ensure_postgres_schema_is_current_is_noop_without_postgres_pool() {
|
||||
let state = AppState::new().expect("state should build");
|
||||
ensure_postgres_schema_is_current(&state)
|
||||
.await
|
||||
.expect("disabled data backend should not block startup");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn explicit_migrate_requires_postgres_url() {
|
||||
let args = test_args();
|
||||
let error = super::run_explicit_migrations(&args)
|
||||
.await
|
||||
.expect_err("missing postgres URL should fail");
|
||||
let message = error.to_string();
|
||||
assert!(message.contains("AETHER_GATEWAY_DATA_POSTGRES_URL or DATABASE_URL"));
|
||||
assert!(message.contains("--migrate"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn explicit_migrate_does_not_depend_on_app_port_validation() {
|
||||
let mut args = test_args();
|
||||
args.app_port = 0;
|
||||
|
||||
let error = super::run_explicit_migrations(&args)
|
||||
.await
|
||||
.expect_err("missing postgres URL should fail before any app port validation");
|
||||
let message = error.to_string();
|
||||
assert!(message.contains("AETHER_GATEWAY_DATA_POSTGRES_URL or DATABASE_URL"));
|
||||
assert!(!message.contains("APP_PORT"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -111,22 +111,24 @@ const MAINTENANCE_DEFAULT_TIMEZONE: &str = "Asia/Shanghai";
|
||||
const DB_MAINTENANCE_TABLES: &[&str] = &["usage", "request_candidates", "audit_logs"];
|
||||
const SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL: &str = r#"
|
||||
SELECT
|
||||
wallet_id,
|
||||
COUNT(id) AS total_requests,
|
||||
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost_usd,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM(cache_creation_input_tokens), 0) AS cache_creation_tokens,
|
||||
COALESCE(SUM(cache_read_input_tokens), 0) AS cache_read_tokens,
|
||||
MIN(finalized_at) AS first_finalized_at,
|
||||
MAX(finalized_at) AS last_finalized_at
|
||||
usage_settlement_snapshots.wallet_id,
|
||||
COUNT(usage.id) AS total_requests,
|
||||
CAST(COALESCE(SUM(usage.total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost_usd,
|
||||
COALESCE(SUM(usage.input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(usage.output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM(usage.cache_creation_input_tokens), 0) AS cache_creation_tokens,
|
||||
COALESCE(SUM(usage.cache_read_input_tokens), 0) AS cache_read_tokens,
|
||||
MIN(COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at)) AS first_finalized_at,
|
||||
MAX(COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at)) AS last_finalized_at
|
||||
FROM usage
|
||||
WHERE wallet_id IS NOT NULL
|
||||
AND billing_status = 'settled'
|
||||
AND total_cost_usd > 0
|
||||
AND finalized_at >= $1
|
||||
AND finalized_at < $2
|
||||
GROUP BY wallet_id
|
||||
JOIN usage_settlement_snapshots
|
||||
ON usage_settlement_snapshots.request_id = usage.request_id
|
||||
WHERE usage_settlement_snapshots.wallet_id IS NOT NULL
|
||||
AND COALESCE(usage_settlement_snapshots.billing_status, usage.billing_status) = 'settled'
|
||||
AND usage.total_cost_usd > 0
|
||||
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) >= $1
|
||||
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) < $2
|
||||
GROUP BY usage_settlement_snapshots.wallet_id
|
||||
"#;
|
||||
const UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL: &str = r#"
|
||||
INSERT INTO wallet_daily_usage_ledgers (
|
||||
@@ -171,11 +173,13 @@ WHERE ledgers.billing_date = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM usage
|
||||
WHERE usage.wallet_id = ledgers.wallet_id
|
||||
AND usage.billing_status = 'settled'
|
||||
JOIN usage_settlement_snapshots
|
||||
ON usage_settlement_snapshots.request_id = usage.request_id
|
||||
WHERE usage_settlement_snapshots.wallet_id = ledgers.wallet_id
|
||||
AND COALESCE(usage_settlement_snapshots.billing_status, usage.billing_status) = 'settled'
|
||||
AND usage.total_cost_usd > 0
|
||||
AND usage.finalized_at >= $3
|
||||
AND usage.finalized_at < $4
|
||||
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) >= $3
|
||||
AND COALESCE(usage_settlement_snapshots.finalized_at, usage.finalized_at) < $4
|
||||
)
|
||||
"#;
|
||||
const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#"
|
||||
@@ -258,7 +262,7 @@ USING doomed
|
||||
WHERE usage_rows.id = doomed.id
|
||||
"#;
|
||||
const SELECT_USAGE_HEADER_BATCH_SQL: &str = r#"
|
||||
SELECT id
|
||||
SELECT id, request_id
|
||||
FROM usage
|
||||
WHERE created_at < $1
|
||||
AND ($2::timestamptz IS NULL OR created_at >= $2)
|
||||
@@ -267,6 +271,17 @@ WHERE created_at < $1
|
||||
OR response_headers IS NOT NULL
|
||||
OR provider_request_headers IS NOT NULL
|
||||
OR client_response_headers IS NOT NULL
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM usage_http_audits
|
||||
WHERE usage_http_audits.request_id = usage.request_id
|
||||
AND (
|
||||
usage_http_audits.request_headers IS NOT NULL
|
||||
OR usage_http_audits.response_headers IS NOT NULL
|
||||
OR usage_http_audits.provider_request_headers IS NOT NULL
|
||||
OR usage_http_audits.client_response_headers IS NOT NULL
|
||||
)
|
||||
)
|
||||
)
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $3
|
||||
@@ -279,8 +294,17 @@ SET request_headers = NULL,
|
||||
client_response_headers = NULL
|
||||
WHERE id = ANY($1)
|
||||
"#;
|
||||
const CLEAR_USAGE_HTTP_AUDIT_HEADERS_SQL: &str = r#"
|
||||
UPDATE usage_http_audits
|
||||
SET request_headers = NULL,
|
||||
response_headers = NULL,
|
||||
provider_request_headers = NULL,
|
||||
client_response_headers = NULL,
|
||||
updated_at = NOW()
|
||||
WHERE request_id = ANY($1)
|
||||
"#;
|
||||
const SELECT_USAGE_STALE_BODY_BATCH_SQL: &str = r#"
|
||||
SELECT id
|
||||
SELECT id, request_id
|
||||
FROM usage
|
||||
WHERE created_at < $1
|
||||
AND ($2::timestamptz IS NULL OR created_at >= $2)
|
||||
@@ -293,6 +317,22 @@ WHERE created_at < $1
|
||||
OR response_body_compressed IS NOT NULL
|
||||
OR provider_request_body_compressed IS NOT NULL
|
||||
OR client_response_body_compressed IS NOT NULL
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM usage_body_blobs
|
||||
WHERE usage_body_blobs.request_id = usage.request_id
|
||||
)
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM usage_http_audits
|
||||
WHERE usage_http_audits.request_id = usage.request_id
|
||||
AND (
|
||||
usage_http_audits.request_body_ref IS NOT NULL
|
||||
OR usage_http_audits.provider_request_body_ref IS NOT NULL
|
||||
OR usage_http_audits.response_body_ref IS NOT NULL
|
||||
OR usage_http_audits.client_response_body_ref IS NOT NULL
|
||||
)
|
||||
)
|
||||
)
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $3
|
||||
@@ -309,35 +349,150 @@ SET request_body = NULL,
|
||||
client_response_body_compressed = NULL
|
||||
WHERE id = ANY($1)
|
||||
"#;
|
||||
const DELETE_USAGE_BODY_BLOBS_SQL: &str = r#"
|
||||
DELETE FROM usage_body_blobs
|
||||
WHERE request_id = ANY($1)
|
||||
"#;
|
||||
const CLEAR_USAGE_HTTP_AUDIT_BODY_REFS_SQL: &str = r#"
|
||||
UPDATE usage_http_audits
|
||||
SET request_body_ref = NULL,
|
||||
provider_request_body_ref = NULL,
|
||||
response_body_ref = NULL,
|
||||
client_response_body_ref = NULL,
|
||||
body_capture_mode = 'none',
|
||||
updated_at = NOW()
|
||||
WHERE request_id = ANY($1)
|
||||
"#;
|
||||
const DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL: &str = r#"
|
||||
DELETE FROM usage_http_audits
|
||||
WHERE request_id = ANY($1)
|
||||
AND request_headers IS NULL
|
||||
AND response_headers IS NULL
|
||||
AND provider_request_headers IS NULL
|
||||
AND client_response_headers IS NULL
|
||||
AND request_body_ref IS NULL
|
||||
AND provider_request_body_ref IS NULL
|
||||
AND response_body_ref IS NULL
|
||||
AND client_response_body_ref IS NULL
|
||||
"#;
|
||||
const SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
request_body,
|
||||
request_body_compressed,
|
||||
response_body,
|
||||
response_body_compressed,
|
||||
provider_request_body,
|
||||
provider_request_body_compressed,
|
||||
client_response_body
|
||||
,
|
||||
client_response_body_compressed
|
||||
FROM usage
|
||||
WHERE created_at < $1
|
||||
AND ($2::timestamptz IS NULL OR created_at >= $2)
|
||||
AND (
|
||||
request_body IS NOT NULL
|
||||
OR request_body_compressed IS NOT NULL
|
||||
OR response_body IS NOT NULL
|
||||
OR response_body_compressed IS NOT NULL
|
||||
OR provider_request_body IS NOT NULL
|
||||
OR provider_request_body_compressed IS NOT NULL
|
||||
OR client_response_body IS NOT NULL
|
||||
OR client_response_body_compressed IS NOT NULL
|
||||
)
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $3
|
||||
"#;
|
||||
const SELECT_USAGE_LEGACY_BODY_REF_METADATA_BATCH_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
request_metadata
|
||||
FROM usage
|
||||
WHERE created_at < $1
|
||||
AND ($2::timestamptz IS NULL OR created_at >= $2)
|
||||
AND request_metadata IS NOT NULL
|
||||
AND (
|
||||
request_metadata::jsonb ? 'request_body_ref'
|
||||
OR request_metadata::jsonb ? 'provider_request_body_ref'
|
||||
OR request_metadata::jsonb ? 'response_body_ref'
|
||||
OR request_metadata::jsonb ? 'client_response_body_ref'
|
||||
)
|
||||
ORDER BY created_at ASC, id ASC
|
||||
LIMIT $3
|
||||
"#;
|
||||
const UPSERT_USAGE_BODY_BLOB_SQL: &str = r#"
|
||||
INSERT INTO usage_body_blobs (
|
||||
body_ref,
|
||||
request_id,
|
||||
body_field,
|
||||
payload_gzip
|
||||
) VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4
|
||||
)
|
||||
ON CONFLICT (body_ref)
|
||||
DO UPDATE SET
|
||||
payload_gzip = EXCLUDED.payload_gzip,
|
||||
updated_at = NOW()
|
||||
"#;
|
||||
const UPDATE_USAGE_REQUEST_METADATA_SQL: &str = r#"
|
||||
UPDATE usage
|
||||
SET request_metadata = $2::json,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
"#;
|
||||
const UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL: &str = r#"
|
||||
INSERT INTO usage_http_audits (
|
||||
request_id,
|
||||
request_body_ref,
|
||||
provider_request_body_ref,
|
||||
response_body_ref,
|
||||
client_response_body_ref,
|
||||
body_capture_mode
|
||||
) VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6
|
||||
)
|
||||
ON CONFLICT (request_id)
|
||||
DO UPDATE SET
|
||||
request_body_ref = COALESCE(EXCLUDED.request_body_ref, usage_http_audits.request_body_ref),
|
||||
provider_request_body_ref = COALESCE(
|
||||
EXCLUDED.provider_request_body_ref,
|
||||
usage_http_audits.provider_request_body_ref
|
||||
),
|
||||
response_body_ref = COALESCE(EXCLUDED.response_body_ref, usage_http_audits.response_body_ref),
|
||||
client_response_body_ref = COALESCE(
|
||||
EXCLUDED.client_response_body_ref,
|
||||
usage_http_audits.client_response_body_ref
|
||||
),
|
||||
body_capture_mode = CASE
|
||||
WHEN EXCLUDED.request_body_ref IS NOT NULL
|
||||
OR EXCLUDED.provider_request_body_ref IS NOT NULL
|
||||
OR EXCLUDED.response_body_ref IS NOT NULL
|
||||
OR EXCLUDED.client_response_body_ref IS NOT NULL
|
||||
THEN EXCLUDED.body_capture_mode
|
||||
ELSE usage_http_audits.body_capture_mode
|
||||
END,
|
||||
updated_at = NOW()
|
||||
"#;
|
||||
const UPDATE_USAGE_BODY_COMPRESSION_SQL: &str = r#"
|
||||
UPDATE usage
|
||||
SET request_body = NULL,
|
||||
response_body = NULL,
|
||||
provider_request_body = NULL,
|
||||
client_response_body = NULL,
|
||||
request_body_compressed = $2,
|
||||
response_body_compressed = $3,
|
||||
provider_request_body_compressed = $4,
|
||||
client_response_body_compressed = $5
|
||||
request_body_compressed = NULL,
|
||||
response_body_compressed = NULL,
|
||||
provider_request_body_compressed = NULL,
|
||||
client_response_body_compressed = NULL
|
||||
WHERE id = $1
|
||||
"#;
|
||||
const SELECT_EXPIRED_ACTIVE_API_KEYS_SQL: &str = r#"
|
||||
@@ -988,7 +1143,8 @@ struct PercentileSummary {
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
struct UsageCleanupSummary {
|
||||
body_compressed: usize,
|
||||
body_externalized: usize,
|
||||
legacy_body_refs_migrated: usize,
|
||||
body_cleaned: usize,
|
||||
header_cleaned: usize,
|
||||
keys_cleaned: usize,
|
||||
@@ -1016,10 +1172,21 @@ struct UsageCleanupWindow {
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
struct UsageBodyCompressionRow {
|
||||
id: String,
|
||||
request_id: String,
|
||||
request_body: Option<Value>,
|
||||
request_body_compressed: Option<Vec<u8>>,
|
||||
response_body: Option<Value>,
|
||||
response_body_compressed: Option<Vec<u8>>,
|
||||
provider_request_body: Option<Value>,
|
||||
provider_request_body_compressed: Option<Vec<u8>>,
|
||||
client_response_body: Option<Value>,
|
||||
client_response_body_compressed: Option<Vec<u8>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct UsageBodyCleanupRow {
|
||||
id: String,
|
||||
request_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
|
||||
@@ -188,7 +188,8 @@ pub(super) async fn run_stats_aggregation_once(
|
||||
|
||||
pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
|
||||
let summary = perform_usage_cleanup_once(data).await?;
|
||||
if summary.body_compressed > 0
|
||||
if summary.body_externalized > 0
|
||||
|| summary.legacy_body_refs_migrated > 0
|
||||
|| summary.body_cleaned > 0
|
||||
|| summary.header_cleaned > 0
|
||||
|| summary.keys_cleaned > 0
|
||||
@@ -198,7 +199,8 @@ pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<()
|
||||
event_name = "usage_cleanup_completed",
|
||||
log_type = "ops",
|
||||
worker = "usage_cleanup",
|
||||
body_compressed = summary.body_compressed,
|
||||
body_externalized = summary.body_externalized,
|
||||
legacy_body_refs_migrated = summary.legacy_body_refs_migrated,
|
||||
body_cleaned = summary.body_cleaned,
|
||||
header_cleaned = summary.header_cleaned,
|
||||
keys_cleaned = summary.keys_cleaned,
|
||||
|
||||
@@ -28,7 +28,8 @@ use super::{
|
||||
stats_hourly_aggregation_target_hour, summarize_postgres_pool, usage_cleanup_settings,
|
||||
usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
|
||||
FailedPendingUsageRow, GatewayDataState, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
|
||||
UsageCleanupSettings, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
|
||||
UsageCleanupSettings, DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL,
|
||||
SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
|
||||
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
||||
};
|
||||
|
||||
@@ -77,6 +78,18 @@ async fn spawn_pool_monitor_worker_skips_when_postgres_unavailable() {
|
||||
assert!(spawn_pool_monitor_worker(Arc::new(GatewayDataState::disabled())).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wallet_daily_usage_queries_use_settlement_snapshots_for_wallet_identity() {
|
||||
assert!(
|
||||
SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL.contains("JOIN usage_settlement_snapshots")
|
||||
);
|
||||
assert!(SELECT_WALLET_DAILY_USAGE_AGGREGATION_ROWS_SQL
|
||||
.contains("usage_settlement_snapshots.wallet_id"));
|
||||
assert!(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL.contains("JOIN usage_settlement_snapshots"));
|
||||
assert!(DELETE_STALE_WALLET_DAILY_USAGE_LEDGERS_SQL
|
||||
.contains("usage_settlement_snapshots.wallet_id = ledgers.wallet_id"));
|
||||
}
|
||||
|
||||
fn sample_connected_proxy_node(
|
||||
node_id: &str,
|
||||
heartbeat_interval: i32,
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
use std::io::Write;
|
||||
|
||||
use aether_data_contracts::repository::usage::{
|
||||
parse_usage_body_ref, usage_body_ref, UsageBodyField,
|
||||
};
|
||||
use aether_data_contracts::DataLayerError;
|
||||
use chrono::{DateTime, Utc};
|
||||
use flate2::{write::GzEncoder, Compression};
|
||||
use serde_json::Value;
|
||||
use serde_json::{Map, Value};
|
||||
use sqlx::Row;
|
||||
use tracing::warn;
|
||||
|
||||
@@ -11,13 +14,17 @@ use crate::data::GatewayDataState;
|
||||
|
||||
use super::{
|
||||
system_config_bool, usage_cleanup_settings, usage_cleanup_window, ExpiredApiKeyRow,
|
||||
UsageBodyCompressionRow, UsageCleanupSummary, CLEAR_USAGE_BODY_FIELDS_SQL,
|
||||
CLEAR_USAGE_HEADER_FIELDS_SQL, DELETE_EXPIRED_API_KEY_SQL, DELETE_OLD_USAGE_RECORDS_SQL,
|
||||
UsageBodyCleanupRow, UsageBodyCompressionRow, UsageCleanupSummary, CLEAR_USAGE_BODY_FIELDS_SQL,
|
||||
CLEAR_USAGE_HEADER_FIELDS_SQL, CLEAR_USAGE_HTTP_AUDIT_BODY_REFS_SQL,
|
||||
CLEAR_USAGE_HTTP_AUDIT_HEADERS_SQL, DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL,
|
||||
DELETE_EXPIRED_API_KEY_SQL, DELETE_OLD_USAGE_RECORDS_SQL, DELETE_USAGE_BODY_BLOBS_SQL,
|
||||
DISABLE_EXPIRED_API_KEY_SQL, EXPIRED_API_KEY_PRE_CLEAN_BATCH_SIZE,
|
||||
NULLIFY_REQUEST_CANDIDATE_API_KEY_BATCH_SQL, NULLIFY_USAGE_API_KEY_BATCH_SQL,
|
||||
SELECT_EXPIRED_ACTIVE_API_KEYS_SQL, SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL,
|
||||
SELECT_USAGE_HEADER_BATCH_SQL, SELECT_USAGE_STALE_BODY_BATCH_SQL,
|
||||
UPDATE_USAGE_BODY_COMPRESSION_SQL,
|
||||
SELECT_USAGE_HEADER_BATCH_SQL, SELECT_USAGE_LEGACY_BODY_REF_METADATA_BATCH_SQL,
|
||||
SELECT_USAGE_STALE_BODY_BATCH_SQL, UPDATE_USAGE_BODY_COMPRESSION_SQL,
|
||||
UPDATE_USAGE_REQUEST_METADATA_SQL, UPSERT_USAGE_BODY_BLOB_SQL,
|
||||
UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL,
|
||||
};
|
||||
|
||||
pub(super) async fn perform_usage_cleanup_once(
|
||||
@@ -41,6 +48,13 @@ pub(super) async fn perform_usage_cleanup_once(
|
||||
Some(window.log_cutoff),
|
||||
)
|
||||
.await?;
|
||||
let legacy_body_refs_migrated = migrate_legacy_usage_body_ref_metadata(
|
||||
&pool,
|
||||
window.detail_cutoff,
|
||||
settings.batch_size,
|
||||
Some(window.compressed_cutoff),
|
||||
)
|
||||
.await?;
|
||||
let body_cleaned = cleanup_usage_stale_body_fields(
|
||||
&pool,
|
||||
window.compressed_cutoff,
|
||||
@@ -48,7 +62,7 @@ pub(super) async fn perform_usage_cleanup_once(
|
||||
Some(window.log_cutoff),
|
||||
)
|
||||
.await?;
|
||||
let body_compressed = compress_usage_body_fields(
|
||||
let body_externalized = compress_usage_body_fields(
|
||||
&pool,
|
||||
window.detail_cutoff,
|
||||
settings.batch_size,
|
||||
@@ -65,7 +79,8 @@ pub(super) async fn perform_usage_cleanup_once(
|
||||
};
|
||||
|
||||
Ok(UsageCleanupSummary {
|
||||
body_compressed,
|
||||
body_externalized,
|
||||
legacy_body_refs_migrated,
|
||||
body_cleaned,
|
||||
header_cleaned,
|
||||
keys_cleaned,
|
||||
@@ -73,6 +88,89 @@ pub(super) async fn perform_usage_cleanup_once(
|
||||
})
|
||||
}
|
||||
|
||||
async fn migrate_legacy_usage_body_ref_metadata(
|
||||
pool: &aether_data::postgres::PostgresPool,
|
||||
cutoff_time: DateTime<Utc>,
|
||||
batch_size: usize,
|
||||
newer_than: Option<DateTime<Utc>>,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
if matches!(newer_than, Some(value) if value >= cutoff_time) {
|
||||
warn!(
|
||||
cutoff_time = %cutoff_time,
|
||||
newer_than = ?newer_than,
|
||||
"gateway usage legacy body-ref migration skipped due to invalid window"
|
||||
);
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let mut total_migrated = 0usize;
|
||||
loop {
|
||||
let rows = sqlx::query(SELECT_USAGE_LEGACY_BODY_REF_METADATA_BATCH_SQL)
|
||||
.bind(cutoff_time)
|
||||
.bind(newer_than)
|
||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
Ok(UsageLegacyBodyRefMetadataRow {
|
||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||
request_id: row
|
||||
.try_get::<String, _>("request_id")
|
||||
.map_err(postgres_error)?,
|
||||
request_metadata: row
|
||||
.try_get::<Option<Value>, _>("request_metadata")
|
||||
.map_err(postgres_error)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, DataLayerError>>()?;
|
||||
if rows.is_empty() {
|
||||
break;
|
||||
}
|
||||
|
||||
let mut batch_migrated = 0usize;
|
||||
for row in rows {
|
||||
let Some(plan) =
|
||||
migrate_legacy_body_ref_metadata_plan(&row.request_id, row.request_metadata)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let mut tx = pool.begin().await.map_err(postgres_error)?;
|
||||
if plan.refs.any_present() {
|
||||
sqlx::query(UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL)
|
||||
.bind(&row.request_id)
|
||||
.bind(plan.refs.request_body_ref.as_deref())
|
||||
.bind(plan.refs.provider_request_body_ref.as_deref())
|
||||
.bind(plan.refs.response_body_ref.as_deref())
|
||||
.bind(plan.refs.client_response_body_ref.as_deref())
|
||||
.bind("ref_backed")
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
}
|
||||
let updated = sqlx::query(UPDATE_USAGE_REQUEST_METADATA_SQL)
|
||||
.bind(&row.id)
|
||||
.bind(plan.request_metadata)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.rows_affected();
|
||||
tx.commit().await.map_err(postgres_error)?;
|
||||
if updated > 0 {
|
||||
batch_migrated += 1;
|
||||
}
|
||||
}
|
||||
|
||||
total_migrated += batch_migrated;
|
||||
if batch_migrated == 0 || batch_migrated < batch_size {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(total_migrated)
|
||||
}
|
||||
|
||||
async fn delete_old_usage_records(
|
||||
pool: &aether_data::postgres::PostgresPool,
|
||||
cutoff_time: DateTime<Utc>,
|
||||
@@ -113,7 +211,7 @@ async fn cleanup_usage_header_fields(
|
||||
|
||||
let mut total_cleaned = 0usize;
|
||||
loop {
|
||||
let ids = sqlx::query(SELECT_USAGE_HEADER_BATCH_SQL)
|
||||
let rows = sqlx::query(SELECT_USAGE_HEADER_BATCH_SQL)
|
||||
.bind(cutoff_time)
|
||||
.bind(newer_than)
|
||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||
@@ -121,11 +219,23 @@ async fn cleanup_usage_header_fields(
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.into_iter()
|
||||
.map(|row| row.try_get::<String, _>("id").map_err(postgres_error))
|
||||
.map(|row| {
|
||||
Ok(UsageBodyCleanupRow {
|
||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||
request_id: row
|
||||
.try_get::<String, _>("request_id")
|
||||
.map_err(postgres_error)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, DataLayerError>>()?;
|
||||
if ids.is_empty() {
|
||||
if rows.is_empty() {
|
||||
break;
|
||||
}
|
||||
let ids = rows.iter().map(|row| row.id.clone()).collect::<Vec<_>>();
|
||||
let request_ids = rows
|
||||
.iter()
|
||||
.map(|row| row.request_id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let cleaned = sqlx::query(CLEAR_USAGE_HEADER_FIELDS_SQL)
|
||||
.bind(ids)
|
||||
@@ -133,6 +243,16 @@ async fn cleanup_usage_header_fields(
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.rows_affected();
|
||||
sqlx::query(CLEAR_USAGE_HTTP_AUDIT_HEADERS_SQL)
|
||||
.bind(&request_ids)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
sqlx::query(DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL)
|
||||
.bind(request_ids)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
let cleaned = usize::try_from(cleaned).unwrap_or(usize::MAX);
|
||||
total_cleaned += cleaned;
|
||||
if cleaned == 0 || cleaned < batch_size {
|
||||
@@ -159,7 +279,7 @@ async fn cleanup_usage_stale_body_fields(
|
||||
|
||||
let mut total_cleaned = 0usize;
|
||||
loop {
|
||||
let ids = sqlx::query(SELECT_USAGE_STALE_BODY_BATCH_SQL)
|
||||
let rows = sqlx::query(SELECT_USAGE_STALE_BODY_BATCH_SQL)
|
||||
.bind(cutoff_time)
|
||||
.bind(newer_than)
|
||||
.bind(i64::try_from(batch_size).unwrap_or(i64::MAX))
|
||||
@@ -167,11 +287,23 @@ async fn cleanup_usage_stale_body_fields(
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.into_iter()
|
||||
.map(|row| row.try_get::<String, _>("id").map_err(postgres_error))
|
||||
.map(|row| {
|
||||
Ok(UsageBodyCleanupRow {
|
||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||
request_id: row
|
||||
.try_get::<String, _>("request_id")
|
||||
.map_err(postgres_error)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, DataLayerError>>()?;
|
||||
if ids.is_empty() {
|
||||
if rows.is_empty() {
|
||||
break;
|
||||
}
|
||||
let ids = rows.iter().map(|row| row.id.clone()).collect::<Vec<_>>();
|
||||
let request_ids = rows
|
||||
.iter()
|
||||
.map(|row| row.request_id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let cleaned = sqlx::query(CLEAR_USAGE_BODY_FIELDS_SQL)
|
||||
.bind(ids)
|
||||
@@ -179,6 +311,21 @@ async fn cleanup_usage_stale_body_fields(
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.rows_affected();
|
||||
sqlx::query(DELETE_USAGE_BODY_BLOBS_SQL)
|
||||
.bind(&request_ids)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
sqlx::query(CLEAR_USAGE_HTTP_AUDIT_BODY_REFS_SQL)
|
||||
.bind(&request_ids)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
sqlx::query(DELETE_EMPTY_USAGE_HTTP_AUDITS_SQL)
|
||||
.bind(request_ids)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
let cleaned = usize::try_from(cleaned).unwrap_or(usize::MAX);
|
||||
total_cleaned += cleaned;
|
||||
if cleaned == 0 || cleaned < batch_size {
|
||||
@@ -218,18 +365,33 @@ async fn compress_usage_body_fields(
|
||||
.map(|row| {
|
||||
Ok(UsageBodyCompressionRow {
|
||||
id: row.try_get::<String, _>("id").map_err(postgres_error)?,
|
||||
request_id: row
|
||||
.try_get::<String, _>("request_id")
|
||||
.map_err(postgres_error)?,
|
||||
request_body: row
|
||||
.try_get::<Option<Value>, _>("request_body")
|
||||
.map_err(postgres_error)?,
|
||||
request_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("request_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
response_body: row
|
||||
.try_get::<Option<Value>, _>("response_body")
|
||||
.map_err(postgres_error)?,
|
||||
response_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("response_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
provider_request_body: row
|
||||
.try_get::<Option<Value>, _>("provider_request_body")
|
||||
.map_err(postgres_error)?,
|
||||
provider_request_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("provider_request_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
client_response_body: row
|
||||
.try_get::<Option<Value>, _>("client_response_body")
|
||||
.map_err(postgres_error)?,
|
||||
client_response_body_compressed: row
|
||||
.try_get::<Option<Vec<u8>>, _>("client_response_body_compressed")
|
||||
.map_err(postgres_error)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, DataLayerError>>()?;
|
||||
@@ -239,18 +401,44 @@ async fn compress_usage_body_fields(
|
||||
|
||||
let mut batch_success = 0usize;
|
||||
for row in rows {
|
||||
let compressed = (
|
||||
compress_usage_json_value(row.request_body.as_ref()),
|
||||
compress_usage_json_value(row.response_body.as_ref()),
|
||||
compress_usage_json_value(row.provider_request_body.as_ref()),
|
||||
compress_usage_json_value(row.client_response_body.as_ref()),
|
||||
);
|
||||
let detached = build_usage_body_externalization(&row)?;
|
||||
if detached.refs.any_present() {
|
||||
let mut tx = pool.begin().await.map_err(postgres_error)?;
|
||||
for blob in &detached.blobs {
|
||||
sqlx::query(UPSERT_USAGE_BODY_BLOB_SQL)
|
||||
.bind(&blob.body_ref)
|
||||
.bind(&row.request_id)
|
||||
.bind(blob.body_field)
|
||||
.bind(&blob.payload_gzip)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
}
|
||||
sqlx::query(UPSERT_USAGE_HTTP_AUDIT_BODY_REFS_SQL)
|
||||
.bind(&row.request_id)
|
||||
.bind(detached.refs.request_body_ref.as_deref())
|
||||
.bind(detached.refs.provider_request_body_ref.as_deref())
|
||||
.bind(detached.refs.response_body_ref.as_deref())
|
||||
.bind(detached.refs.client_response_body_ref.as_deref())
|
||||
.bind("ref_backed")
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(postgres_error)?;
|
||||
let updated = sqlx::query(UPDATE_USAGE_BODY_COMPRESSION_SQL)
|
||||
.bind(&row.id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
.rows_affected();
|
||||
tx.commit().await.map_err(postgres_error)?;
|
||||
if updated > 0 {
|
||||
batch_success += 1;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let updated = sqlx::query(UPDATE_USAGE_BODY_COMPRESSION_SQL)
|
||||
.bind(row.id)
|
||||
.bind(compressed.0)
|
||||
.bind(compressed.1)
|
||||
.bind(compressed.2)
|
||||
.bind(compressed.3)
|
||||
.bind(&row.id)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(postgres_error)?
|
||||
@@ -276,12 +464,174 @@ async fn compress_usage_body_fields(
|
||||
Ok(total_compressed)
|
||||
}
|
||||
|
||||
fn compress_usage_json_value(value: Option<&Value>) -> Option<Vec<u8>> {
|
||||
let value = value?;
|
||||
let bytes = serde_json::to_vec(value).ok()?;
|
||||
fn compress_usage_json_value(value: &Value) -> Result<Vec<u8>, DataLayerError> {
|
||||
let bytes = serde_json::to_vec(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("failed to serialize usage json for gzip: {err}"))
|
||||
})?;
|
||||
let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6));
|
||||
encoder.write_all(&bytes).ok()?;
|
||||
encoder.finish().ok()
|
||||
encoder.write_all(&bytes).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("failed to gzip usage json: {err}"))
|
||||
})?;
|
||||
encoder.finish().map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("failed to finish gzipped usage json: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct UsageDetachedBodyBlobWrite {
|
||||
body_ref: String,
|
||||
body_field: &'static str,
|
||||
payload_gzip: Vec<u8>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
struct UsageDetachedBodyRefs {
|
||||
request_body_ref: Option<String>,
|
||||
provider_request_body_ref: Option<String>,
|
||||
response_body_ref: Option<String>,
|
||||
client_response_body_ref: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
struct UsageLegacyBodyRefMetadataRow {
|
||||
id: String,
|
||||
request_id: String,
|
||||
request_metadata: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
struct UsageLegacyBodyRefMigrationPlan {
|
||||
refs: UsageDetachedBodyRefs,
|
||||
request_metadata: Option<Value>,
|
||||
}
|
||||
|
||||
impl UsageDetachedBodyRefs {
|
||||
fn any_present(&self) -> bool {
|
||||
self.request_body_ref.is_some()
|
||||
|| self.provider_request_body_ref.is_some()
|
||||
|| self.response_body_ref.is_some()
|
||||
|| self.client_response_body_ref.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
struct UsageBodyExternalizationPlan {
|
||||
blobs: Vec<UsageDetachedBodyBlobWrite>,
|
||||
refs: UsageDetachedBodyRefs,
|
||||
}
|
||||
|
||||
fn migrate_legacy_body_ref_metadata_plan(
|
||||
request_id: &str,
|
||||
request_metadata: Option<Value>,
|
||||
) -> Option<UsageLegacyBodyRefMigrationPlan> {
|
||||
let mut metadata = match request_metadata {
|
||||
Some(Value::Object(object)) => object,
|
||||
_ => return None,
|
||||
};
|
||||
|
||||
let mut refs = UsageDetachedBodyRefs::default();
|
||||
let mut removed_any = false;
|
||||
for field in [
|
||||
UsageBodyField::RequestBody,
|
||||
UsageBodyField::ProviderRequestBody,
|
||||
UsageBodyField::ResponseBody,
|
||||
UsageBodyField::ClientResponseBody,
|
||||
] {
|
||||
let key = field.as_ref_key();
|
||||
let Some(value) = metadata.remove(key) else {
|
||||
continue;
|
||||
};
|
||||
removed_any = true;
|
||||
let parsed = value
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(parse_usage_body_ref)
|
||||
.filter(|(parsed_request_id, parsed_field)| {
|
||||
parsed_request_id == request_id && *parsed_field == field
|
||||
})
|
||||
.map(|(parsed_request_id, parsed_field)| {
|
||||
usage_body_ref(&parsed_request_id, parsed_field)
|
||||
});
|
||||
match field {
|
||||
UsageBodyField::RequestBody => refs.request_body_ref = parsed,
|
||||
UsageBodyField::ProviderRequestBody => refs.provider_request_body_ref = parsed,
|
||||
UsageBodyField::ResponseBody => refs.response_body_ref = parsed,
|
||||
UsageBodyField::ClientResponseBody => refs.client_response_body_ref = parsed,
|
||||
}
|
||||
}
|
||||
|
||||
if !removed_any {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(UsageLegacyBodyRefMigrationPlan {
|
||||
refs,
|
||||
request_metadata: (!metadata.is_empty()).then_some(Value::Object(metadata)),
|
||||
})
|
||||
}
|
||||
|
||||
fn build_usage_body_externalization(
|
||||
row: &UsageBodyCompressionRow,
|
||||
) -> Result<UsageBodyExternalizationPlan, DataLayerError> {
|
||||
let mut plan = UsageBodyExternalizationPlan::default();
|
||||
maybe_externalize_usage_body_field(
|
||||
&mut plan,
|
||||
&row.request_id,
|
||||
UsageBodyField::RequestBody,
|
||||
row.request_body.as_ref(),
|
||||
row.request_body_compressed.as_deref(),
|
||||
)?;
|
||||
maybe_externalize_usage_body_field(
|
||||
&mut plan,
|
||||
&row.request_id,
|
||||
UsageBodyField::ProviderRequestBody,
|
||||
row.provider_request_body.as_ref(),
|
||||
row.provider_request_body_compressed.as_deref(),
|
||||
)?;
|
||||
maybe_externalize_usage_body_field(
|
||||
&mut plan,
|
||||
&row.request_id,
|
||||
UsageBodyField::ResponseBody,
|
||||
row.response_body.as_ref(),
|
||||
row.response_body_compressed.as_deref(),
|
||||
)?;
|
||||
maybe_externalize_usage_body_field(
|
||||
&mut plan,
|
||||
&row.request_id,
|
||||
UsageBodyField::ClientResponseBody,
|
||||
row.client_response_body.as_ref(),
|
||||
row.client_response_body_compressed.as_deref(),
|
||||
)?;
|
||||
Ok(plan)
|
||||
}
|
||||
|
||||
fn maybe_externalize_usage_body_field(
|
||||
plan: &mut UsageBodyExternalizationPlan,
|
||||
request_id: &str,
|
||||
field: UsageBodyField,
|
||||
inline_body: Option<&Value>,
|
||||
compressed_body: Option<&[u8]>,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let Some(payload_gzip) = (match inline_body {
|
||||
Some(value) => Some(compress_usage_json_value(value)?),
|
||||
None => compressed_body.map(|value| value.to_vec()),
|
||||
}) else {
|
||||
return Ok(());
|
||||
};
|
||||
let body_ref = usage_body_ref(request_id, field);
|
||||
plan.blobs.push(UsageDetachedBodyBlobWrite {
|
||||
body_ref: body_ref.clone(),
|
||||
body_field: field.as_storage_field(),
|
||||
payload_gzip,
|
||||
});
|
||||
match field {
|
||||
UsageBodyField::RequestBody => plan.refs.request_body_ref = Some(body_ref),
|
||||
UsageBodyField::ProviderRequestBody => plan.refs.provider_request_body_ref = Some(body_ref),
|
||||
UsageBodyField::ResponseBody => plan.refs.response_body_ref = Some(body_ref),
|
||||
UsageBodyField::ClientResponseBody => plan.refs.client_response_body_ref = Some(body_ref),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cleanup_expired_api_keys(
|
||||
@@ -375,3 +725,137 @@ async fn nullify_expired_api_key_candidate_refs(
|
||||
fn postgres_error(error: sqlx::Error) -> DataLayerError {
|
||||
DataLayerError::postgres(error)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::io::Read;
|
||||
|
||||
use flate2::read::GzDecoder;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
build_usage_body_externalization, compress_usage_json_value,
|
||||
migrate_legacy_body_ref_metadata_plan, UsageBodyCompressionRow,
|
||||
};
|
||||
|
||||
fn inflate_json(bytes: &[u8]) -> serde_json::Value {
|
||||
let mut decoder = GzDecoder::new(bytes);
|
||||
let mut decoded = Vec::new();
|
||||
decoder
|
||||
.read_to_end(&mut decoded)
|
||||
.expect("gzip should decode");
|
||||
serde_json::from_slice(&decoded).expect("json should decode")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_body_externalization_moves_inline_json_into_ref_backed_blobs() {
|
||||
let row = UsageBodyCompressionRow {
|
||||
id: "usage-1".to_string(),
|
||||
request_id: "req-1".to_string(),
|
||||
request_body: Some(json!({"hello": "world"})),
|
||||
request_body_compressed: None,
|
||||
response_body: None,
|
||||
response_body_compressed: None,
|
||||
provider_request_body: Some(json!({"provider": true})),
|
||||
provider_request_body_compressed: None,
|
||||
client_response_body: None,
|
||||
client_response_body_compressed: None,
|
||||
};
|
||||
|
||||
let plan = build_usage_body_externalization(&row).expect("plan should build");
|
||||
|
||||
assert_eq!(plan.blobs.len(), 2);
|
||||
assert_eq!(
|
||||
plan.refs.request_body_ref.as_deref(),
|
||||
Some("usage://request/req-1/request_body")
|
||||
);
|
||||
assert_eq!(
|
||||
plan.refs.provider_request_body_ref.as_deref(),
|
||||
Some("usage://request/req-1/provider_request_body")
|
||||
);
|
||||
assert_eq!(
|
||||
inflate_json(&plan.blobs[0].payload_gzip),
|
||||
json!({"hello": "world"})
|
||||
);
|
||||
assert_eq!(
|
||||
inflate_json(&plan.blobs[1].payload_gzip),
|
||||
json!({"provider": true})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_body_externalization_reuses_existing_compressed_payloads() {
|
||||
let compressed = compress_usage_json_value(&json!({"legacy": true}))
|
||||
.expect("compressed payload should build");
|
||||
let row = UsageBodyCompressionRow {
|
||||
id: "usage-1".to_string(),
|
||||
request_id: "req-legacy".to_string(),
|
||||
request_body: None,
|
||||
request_body_compressed: Some(compressed.clone()),
|
||||
response_body: None,
|
||||
response_body_compressed: None,
|
||||
provider_request_body: None,
|
||||
provider_request_body_compressed: None,
|
||||
client_response_body: None,
|
||||
client_response_body_compressed: None,
|
||||
};
|
||||
|
||||
let plan = build_usage_body_externalization(&row).expect("plan should build");
|
||||
|
||||
assert_eq!(plan.blobs.len(), 1);
|
||||
assert_eq!(plan.blobs[0].payload_gzip, compressed);
|
||||
assert_eq!(
|
||||
plan.refs.request_body_ref.as_deref(),
|
||||
Some("usage://request/req-legacy/request_body")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn legacy_body_ref_metadata_migration_moves_matching_refs_and_strips_keys() {
|
||||
let plan = migrate_legacy_body_ref_metadata_plan(
|
||||
"req-1",
|
||||
Some(json!({
|
||||
"trace_id": "trace-1",
|
||||
"request_body_ref": "usage://request/req-1/request_body",
|
||||
"response_body_ref": "usage://request/req-1/response_body"
|
||||
})),
|
||||
)
|
||||
.expect("migration plan should exist");
|
||||
|
||||
assert_eq!(
|
||||
plan.refs.request_body_ref.as_deref(),
|
||||
Some("usage://request/req-1/request_body")
|
||||
);
|
||||
assert_eq!(
|
||||
plan.refs.response_body_ref.as_deref(),
|
||||
Some("usage://request/req-1/response_body")
|
||||
);
|
||||
assert_eq!(
|
||||
plan.request_metadata,
|
||||
Some(json!({
|
||||
"trace_id": "trace-1"
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn legacy_body_ref_metadata_migration_strips_invalid_and_cross_request_refs() {
|
||||
let plan = migrate_legacy_body_ref_metadata_plan(
|
||||
"req-1",
|
||||
Some(json!({
|
||||
"request_body_ref": "blob://legacy-request",
|
||||
"provider_request_body_ref": "usage://request/req-other/provider_request_body",
|
||||
"candidate_index": 2
|
||||
})),
|
||||
)
|
||||
.expect("migration plan should exist");
|
||||
|
||||
assert!(!plan.refs.any_present());
|
||||
assert_eq!(
|
||||
plan.request_metadata,
|
||||
Some(json!({
|
||||
"candidate_index": 2
|
||||
}))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -67,33 +67,18 @@ pub(crate) async fn access_log_middleware(mut request: Request<Body>, next: Next
|
||||
HeaderValue::from_str(&trace_id).expect("trace id should be a valid header value"),
|
||||
);
|
||||
}
|
||||
if should_downgrade_access_log(&method, &path) {
|
||||
trace!(
|
||||
event_name = "http_request_started",
|
||||
log_type = "access",
|
||||
status = "started",
|
||||
trace_id = %trace_id,
|
||||
request_id = "-",
|
||||
method = %method,
|
||||
path = %path,
|
||||
route_class = "pending",
|
||||
execution_path = "pending",
|
||||
"gateway request started"
|
||||
);
|
||||
} else {
|
||||
info!(
|
||||
event_name = "http_request_started",
|
||||
log_type = "access",
|
||||
status = "started",
|
||||
trace_id = %trace_id,
|
||||
request_id = "-",
|
||||
method = %method,
|
||||
path = %path,
|
||||
route_class = "pending",
|
||||
execution_path = "pending",
|
||||
"gateway request started"
|
||||
);
|
||||
}
|
||||
trace!(
|
||||
event_name = "http_request_started",
|
||||
log_type = "access",
|
||||
status = "started",
|
||||
trace_id = %trace_id,
|
||||
request_id = "-",
|
||||
method = %method,
|
||||
path = %path,
|
||||
route_class = "pending",
|
||||
execution_path = "pending",
|
||||
"gateway request started"
|
||||
);
|
||||
let mut response = next.run(request).await;
|
||||
if !response.headers().contains_key(TRACE_ID_HEADER) {
|
||||
response.headers_mut().insert(
|
||||
@@ -228,7 +213,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn access_log_emits_started_and_completed_events() {
|
||||
async fn access_log_emits_completed_events_by_default() {
|
||||
let writer = SharedBuffer::default();
|
||||
let subscriber = tracing_subscriber::registry().with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
@@ -236,7 +221,8 @@ mod tests {
|
||||
.flatten_event(true)
|
||||
.with_current_span(false)
|
||||
.with_span_list(false)
|
||||
.with_writer(writer.clone()),
|
||||
.with_writer(writer.clone())
|
||||
.with_filter(LevelFilter::INFO),
|
||||
);
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
let _guard = tracing::dispatcher::set_default(&dispatch);
|
||||
@@ -276,15 +262,13 @@ mod tests {
|
||||
assert!(response.headers().contains_key(TRACE_ID_HEADER));
|
||||
|
||||
let logs = writer.lines();
|
||||
assert_eq!(logs.len(), 2);
|
||||
assert_eq!(logs[0]["event_name"], "http_request_started");
|
||||
assert_eq!(logs[0]["status"], "started");
|
||||
assert_eq!(logs[1]["event_name"], "http_request_completed");
|
||||
assert_eq!(logs[1]["status"], "completed");
|
||||
assert_eq!(logs[1]["status_code"], 200);
|
||||
assert_eq!(logs[1]["request_id"], "req-123");
|
||||
assert_eq!(logs[1]["route_class"], "local");
|
||||
assert_eq!(logs[1]["execution_path"], "local_route");
|
||||
assert_eq!(logs.len(), 1);
|
||||
assert_eq!(logs[0]["event_name"], "http_request_completed");
|
||||
assert_eq!(logs[0]["status"], "completed");
|
||||
assert_eq!(logs[0]["status_code"], 200);
|
||||
assert_eq!(logs[0]["request_id"], "req-123");
|
||||
assert_eq!(logs[0]["route_class"], "local");
|
||||
assert_eq!(logs[0]["execution_path"], "local_route");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -343,7 +327,8 @@ mod tests {
|
||||
.flatten_event(true)
|
||||
.with_current_span(false)
|
||||
.with_span_list(false)
|
||||
.with_writer(writer.clone()),
|
||||
.with_writer(writer.clone())
|
||||
.with_filter(LevelFilter::INFO),
|
||||
);
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
let _guard = tracing::dispatcher::set_default(&dispatch);
|
||||
@@ -383,11 +368,11 @@ mod tests {
|
||||
.expect("request should succeed");
|
||||
|
||||
let logs = writer.lines();
|
||||
assert_eq!(logs[1]["request_id"], "d07e1e94");
|
||||
assert_eq!(logs[0]["request_id"], "d07e1e94");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn access_log_emits_started_and_failed_events_for_server_errors() {
|
||||
async fn access_log_emits_failed_events_by_default_for_server_errors() {
|
||||
let writer = SharedBuffer::default();
|
||||
let subscriber = tracing_subscriber::registry().with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
@@ -395,7 +380,8 @@ mod tests {
|
||||
.flatten_event(true)
|
||||
.with_current_span(false)
|
||||
.with_span_list(false)
|
||||
.with_writer(writer.clone()),
|
||||
.with_writer(writer.clone())
|
||||
.with_filter(LevelFilter::INFO),
|
||||
);
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
let _guard = tracing::dispatcher::set_default(&dispatch);
|
||||
@@ -425,13 +411,12 @@ mod tests {
|
||||
.expect("request should succeed");
|
||||
|
||||
let logs = writer.lines();
|
||||
assert_eq!(logs.len(), 2);
|
||||
assert_eq!(logs[0]["event_name"], "http_request_started");
|
||||
assert_eq!(logs[1]["event_name"], "http_request_failed");
|
||||
assert_eq!(logs[1]["status"], "failed");
|
||||
assert_eq!(logs[1]["status_code"], 502);
|
||||
assert_eq!(logs[1]["route_class"], "passthrough");
|
||||
assert_eq!(logs[1]["execution_path"], "execution_runtime_sync");
|
||||
assert_eq!(logs.len(), 1);
|
||||
assert_eq!(logs[0]["event_name"], "http_request_failed");
|
||||
assert_eq!(logs[0]["status"], "failed");
|
||||
assert_eq!(logs[0]["status_code"], 502);
|
||||
assert_eq!(logs[0]["route_class"], "passthrough");
|
||||
assert_eq!(logs[0]["execution_path"], "execution_runtime_sync");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -443,7 +428,8 @@ mod tests {
|
||||
.flatten_event(true)
|
||||
.with_current_span(false)
|
||||
.with_span_list(false)
|
||||
.with_writer(writer.clone()),
|
||||
.with_writer(writer.clone())
|
||||
.with_filter(LevelFilter::INFO),
|
||||
);
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
let _guard = tracing::dispatcher::set_default(&dispatch);
|
||||
@@ -473,13 +459,12 @@ mod tests {
|
||||
.expect("request should succeed");
|
||||
|
||||
let logs = writer.lines();
|
||||
assert_eq!(logs.len(), 2);
|
||||
assert_eq!(logs[0]["event_name"], "http_request_started");
|
||||
assert_eq!(logs[1]["event_name"], "http_request_completed");
|
||||
assert_eq!(logs[1]["status"], "completed");
|
||||
assert_eq!(logs[1]["status_code"], 401);
|
||||
assert_eq!(logs[1]["route_class"], "auth");
|
||||
assert_eq!(logs[1]["execution_path"], "local_auth_denied");
|
||||
assert_eq!(logs.len(), 1);
|
||||
assert_eq!(logs[0]["event_name"], "http_request_completed");
|
||||
assert_eq!(logs[0]["status"], "completed");
|
||||
assert_eq!(logs[0]["status_code"], 401);
|
||||
assert_eq!(logs[0]["route_class"], "auth");
|
||||
assert_eq!(logs[0]["execution_path"], "local_auth_denied");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -491,7 +476,8 @@ mod tests {
|
||||
.flatten_event(true)
|
||||
.with_current_span(false)
|
||||
.with_span_list(false)
|
||||
.with_writer(writer.clone()),
|
||||
.with_writer(writer.clone())
|
||||
.with_filter(LevelFilter::INFO),
|
||||
);
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
let _guard = tracing::dispatcher::set_default(&dispatch);
|
||||
@@ -528,13 +514,12 @@ mod tests {
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let logs = writer.lines();
|
||||
assert_eq!(logs.len(), 2);
|
||||
assert_eq!(logs[0]["event_name"], "http_request_started");
|
||||
assert_eq!(logs[1]["event_name"], "http_request_completed");
|
||||
assert_eq!(logs[1]["status_code"], 200);
|
||||
assert_eq!(logs[1]["request_id"], "req-stream");
|
||||
assert_eq!(logs[1]["route_class"], "ai_public");
|
||||
assert_eq!(logs[1]["execution_path"], "execution_runtime_stream");
|
||||
assert_eq!(logs.len(), 1);
|
||||
assert_eq!(logs[0]["event_name"], "http_request_completed");
|
||||
assert_eq!(logs[0]["status_code"], 200);
|
||||
assert_eq!(logs[0]["request_id"], "req-stream");
|
||||
assert_eq!(logs[0]["route_class"], "ai_public");
|
||||
assert_eq!(logs[0]["execution_path"], "execution_runtime_stream");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -252,6 +252,28 @@ impl AppState {
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub async fn pending_postgres_migrations(
|
||||
&self,
|
||||
) -> Result<Option<Vec<aether_data::migrate::PendingMigrationInfo>>, sqlx::migrate::MigrateError>
|
||||
{
|
||||
let Some(pool) = self.postgres_pool() else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(aether_data::migrate::pending_migrations(&pool).await?))
|
||||
}
|
||||
|
||||
pub async fn prepare_postgres_for_startup(
|
||||
&self,
|
||||
) -> Result<Option<Vec<aether_data::migrate::PendingMigrationInfo>>, sqlx::migrate::MigrateError>
|
||||
{
|
||||
let Some(pool) = self.postgres_pool() else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(
|
||||
aether_data::migrate::prepare_database_for_startup(&pool).await?,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn with_video_task_poller_config(mut self, interval: Duration, batch_size: usize) -> Self {
|
||||
self.video_task_poller = Some(VideoTaskPollerConfig {
|
||||
interval,
|
||||
@@ -559,10 +581,17 @@ impl AppState {
|
||||
trace_id: &str,
|
||||
diagnostic: LocalExecutionRuntimeMissDiagnostic,
|
||||
) {
|
||||
self.local_execution_runtime_miss_diagnostics
|
||||
let mut diagnostics = self
|
||||
.local_execution_runtime_miss_diagnostics
|
||||
.lock()
|
||||
.expect("local execution runtime miss diagnostics should lock")
|
||||
.insert(trace_id.to_string(), diagnostic);
|
||||
.expect("local execution runtime miss diagnostics should lock");
|
||||
if diagnostics
|
||||
.get(trace_id)
|
||||
.is_some_and(|existing| should_preserve_runtime_miss_diagnostic(existing, &diagnostic))
|
||||
{
|
||||
return;
|
||||
}
|
||||
diagnostics.insert(trace_id.to_string(), diagnostic);
|
||||
}
|
||||
|
||||
pub(crate) fn mutate_local_execution_runtime_miss_diagnostic<F>(
|
||||
@@ -581,6 +610,17 @@ impl AppState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn local_execution_runtime_miss_diagnostic_has_candidate_signal(
|
||||
&self,
|
||||
trace_id: &str,
|
||||
) -> bool {
|
||||
self.local_execution_runtime_miss_diagnostics
|
||||
.lock()
|
||||
.expect("local execution runtime miss diagnostics should lock")
|
||||
.get(trace_id)
|
||||
.is_some_and(runtime_miss_diagnostic_has_candidate_signal)
|
||||
}
|
||||
|
||||
pub(crate) fn take_local_execution_runtime_miss_diagnostic(
|
||||
&self,
|
||||
trace_id: &str,
|
||||
@@ -735,3 +775,19 @@ impl AppState {
|
||||
tasks
|
||||
}
|
||||
}
|
||||
|
||||
fn should_preserve_runtime_miss_diagnostic(
|
||||
existing: &LocalExecutionRuntimeMissDiagnostic,
|
||||
next: &LocalExecutionRuntimeMissDiagnostic,
|
||||
) -> bool {
|
||||
runtime_miss_diagnostic_has_candidate_signal(existing)
|
||||
&& !runtime_miss_diagnostic_has_candidate_signal(next)
|
||||
}
|
||||
|
||||
fn runtime_miss_diagnostic_has_candidate_signal(
|
||||
diagnostic: &LocalExecutionRuntimeMissDiagnostic,
|
||||
) -> bool {
|
||||
diagnostic.candidate_count.unwrap_or(0) > 0
|
||||
|| diagnostic.skipped_candidate_count.unwrap_or(0) > 0
|
||||
|| !diagnostic.skip_reasons.is_empty()
|
||||
}
|
||||
|
||||
@@ -1,8 +1,108 @@
|
||||
use crate::AppState;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use aether_data::repository::auth::{AuthApiKeyLookupKey, ResolvedAuthApiKeySnapshotReader};
|
||||
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::super::super::{AUTH_API_KEY_LAST_USED_MAX_ENTRIES, AUTH_API_KEY_LAST_USED_TTL};
|
||||
|
||||
impl AppState {
|
||||
pub(crate) async fn resolve_auth_api_key_snapshots_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<aether_data::repository::auth::StoredAuthApiKeySnapshot>, GatewayError> {
|
||||
if !self.has_auth_api_key_data_reader() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let api_key_ids = api_key_ids
|
||||
.iter()
|
||||
.map(|value| value.trim())
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
if api_key_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut snapshots = self
|
||||
.data
|
||||
.list_auth_api_key_snapshots_by_ids(&api_key_ids)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
.into_iter()
|
||||
.map(|snapshot| (snapshot.api_key_id.clone(), snapshot))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
|
||||
for api_key_id in &api_key_ids {
|
||||
if snapshots.contains_key(api_key_id) {
|
||||
continue;
|
||||
}
|
||||
let snapshot = self
|
||||
.data
|
||||
.find_stored_auth_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId(api_key_id))
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if let Some(snapshot) = snapshot {
|
||||
snapshots.insert(api_key_id.clone(), snapshot);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(snapshots.into_values().collect())
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_auth_api_key_names_by_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<BTreeMap<String, String>, GatewayError> {
|
||||
if !self.has_auth_api_key_data_reader() {
|
||||
return Ok(BTreeMap::new());
|
||||
}
|
||||
|
||||
let api_key_ids = api_key_ids
|
||||
.iter()
|
||||
.map(|value| value.trim())
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
if api_key_ids.is_empty() {
|
||||
return Ok(BTreeMap::new());
|
||||
}
|
||||
|
||||
let mut names = self
|
||||
.data
|
||||
.list_auth_api_key_snapshots_by_ids(&api_key_ids)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||
.into_iter()
|
||||
.filter_map(|snapshot| {
|
||||
snapshot
|
||||
.api_key_name
|
||||
.map(|name| (snapshot.api_key_id, name))
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
|
||||
for api_key_id in &api_key_ids {
|
||||
if names.contains_key(api_key_id) {
|
||||
continue;
|
||||
}
|
||||
let snapshot = self
|
||||
.data
|
||||
.find_stored_auth_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId(api_key_id))
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if let Some(name) = snapshot.and_then(|snapshot| snapshot.api_key_name) {
|
||||
names.insert(api_key_id.clone(), name);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(names)
|
||||
}
|
||||
|
||||
pub(crate) async fn touch_auth_api_key_last_used_best_effort(&self, api_key_id: &str) {
|
||||
let api_key_id = api_key_id.trim();
|
||||
if api_key_id.is_empty() || !self.data.has_auth_api_key_writer() {
|
||||
@@ -24,3 +124,192 @@ impl AppState {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::auth::{
|
||||
AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository,
|
||||
InMemoryAuthApiKeySnapshotRepository, StandaloneApiKeyExportListQuery,
|
||||
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::AppState;
|
||||
|
||||
#[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
|
||||
}
|
||||
}
|
||||
|
||||
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("[email protected]".to_string()),
|
||||
"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 api key snapshot should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_auth_api_key_names_by_ids_falls_back_to_single_lookup_for_missing_list_rows() {
|
||||
let repository = Arc::new(PartialListAuthApiKeyRepository {
|
||||
lookup: InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
None,
|
||||
sample_usage_auth_snapshot("key-1", "user-1", "fresh-default"),
|
||||
)]),
|
||||
});
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository);
|
||||
|
||||
let names = state
|
||||
.resolve_auth_api_key_names_by_ids(&["key-1".to_string()])
|
||||
.await
|
||||
.expect("api key name resolution should succeed");
|
||||
|
||||
assert_eq!(
|
||||
names,
|
||||
BTreeMap::from([("key-1".to_string(), "fresh-default".to_string())])
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_auth_api_key_snapshots_by_ids_falls_back_to_single_lookup_for_missing_list_rows(
|
||||
) {
|
||||
let repository = Arc::new(PartialListAuthApiKeyRepository {
|
||||
lookup: InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
None,
|
||||
sample_usage_auth_snapshot("key-1", "user-1", "fresh-default"),
|
||||
)]),
|
||||
});
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository);
|
||||
|
||||
let snapshots = state
|
||||
.resolve_auth_api_key_snapshots_by_ids(&["key-1".to_string()])
|
||||
.await
|
||||
.expect("api key snapshot resolution should succeed");
|
||||
|
||||
assert_eq!(snapshots.len(), 1);
|
||||
assert_eq!(snapshots[0].api_key_id, "key-1");
|
||||
assert_eq!(snapshots[0].api_key_name.as_deref(), Some("fresh-default"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,48 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
impl AppState {
|
||||
pub(crate) async fn resolve_auth_user_summaries_by_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<BTreeMap<String, aether_data::repository::users::StoredUserSummary>, GatewayError>
|
||||
{
|
||||
let user_ids = user_ids
|
||||
.iter()
|
||||
.map(|value| value.trim())
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
if user_ids.is_empty() {
|
||||
return Ok(BTreeMap::new());
|
||||
}
|
||||
|
||||
let mut users = BTreeMap::new();
|
||||
if self.has_user_data_reader() {
|
||||
for user in self.list_users_by_ids(&user_ids).await? {
|
||||
users.insert(user.id.clone(), user);
|
||||
}
|
||||
}
|
||||
|
||||
for user_id in &user_ids {
|
||||
if users.contains_key(user_id) {
|
||||
continue;
|
||||
}
|
||||
let Some(user) = self.find_user_auth_by_id(user_id).await? else {
|
||||
continue;
|
||||
};
|
||||
let summary = user
|
||||
.to_summary()
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
users.insert(summary.id.clone(), summary);
|
||||
}
|
||||
|
||||
Ok(users)
|
||||
}
|
||||
|
||||
pub(crate) async fn find_user_auth_by_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
|
||||
@@ -243,6 +243,56 @@ impl AppState {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn initialize_auth_api_key_wallet(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
initial_gift_usd: f64,
|
||||
unlimited: bool,
|
||||
) -> Result<Option<aether_data::repository::wallet::StoredWalletSnapshot>, GatewayError> {
|
||||
#[cfg(test)]
|
||||
if let Some(store) = self.auth_wallet_store.as_ref() {
|
||||
let gift_balance = if unlimited {
|
||||
0.0
|
||||
} else {
|
||||
initial_gift_usd.max(0.0)
|
||||
};
|
||||
let now_unix_secs = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs() as i64;
|
||||
let wallet = aether_data::repository::wallet::StoredWalletSnapshot::new(
|
||||
uuid::Uuid::new_v4().to_string(),
|
||||
None,
|
||||
Some(api_key_id.to_string()),
|
||||
0.0,
|
||||
gift_balance,
|
||||
if unlimited {
|
||||
"unlimited".to_string()
|
||||
} else {
|
||||
"finite".to_string()
|
||||
},
|
||||
"USD".to_string(),
|
||||
"active".to_string(),
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
gift_balance,
|
||||
now_unix_secs,
|
||||
)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
store
|
||||
.lock()
|
||||
.expect("auth wallet store should lock")
|
||||
.insert(wallet.id.clone(), wallet.clone());
|
||||
return Ok(Some(wallet));
|
||||
}
|
||||
|
||||
self.data
|
||||
.initialize_auth_api_key_wallet(api_key_id, initial_gift_usd, unlimited)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn update_auth_user_wallet_limit_mode(
|
||||
&self,
|
||||
user_id: &str,
|
||||
@@ -268,4 +318,144 @@ impl AppState {
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn update_auth_api_key_wallet_limit_mode(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
limit_mode: &str,
|
||||
) -> Result<Option<aether_data::repository::wallet::StoredWalletSnapshot>, GatewayError> {
|
||||
#[cfg(test)]
|
||||
if let Some(store) = self.auth_wallet_store.as_ref() {
|
||||
let mut guard = store.lock().expect("auth wallet store should lock");
|
||||
let Some((wallet_id, wallet)) = guard
|
||||
.iter_mut()
|
||||
.find(|(_, wallet)| wallet.api_key_id.as_deref() == Some(api_key_id))
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let _ = wallet_id;
|
||||
wallet.limit_mode = limit_mode.to_string();
|
||||
wallet.updated_at_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
|
||||
return Ok(Some(wallet.clone()));
|
||||
}
|
||||
|
||||
self.data
|
||||
.update_auth_api_key_wallet_limit_mode(api_key_id, limit_mode)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn update_auth_user_wallet_snapshot(
|
||||
&self,
|
||||
user_id: &str,
|
||||
balance: f64,
|
||||
gift_balance: f64,
|
||||
limit_mode: &str,
|
||||
currency: &str,
|
||||
status: &str,
|
||||
total_recharged: f64,
|
||||
total_consumed: f64,
|
||||
total_refunded: f64,
|
||||
total_adjusted: f64,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Result<Option<aether_data::repository::wallet::StoredWalletSnapshot>, GatewayError> {
|
||||
#[cfg(test)]
|
||||
if let Some(store) = self.auth_wallet_store.as_ref() {
|
||||
let mut guard = store.lock().expect("auth wallet store should lock");
|
||||
let Some((_, wallet)) = guard
|
||||
.iter_mut()
|
||||
.find(|(_, wallet)| wallet.user_id.as_deref() == Some(user_id))
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
wallet.balance = balance;
|
||||
wallet.gift_balance = gift_balance;
|
||||
wallet.limit_mode = limit_mode.to_string();
|
||||
wallet.currency = currency.to_string();
|
||||
wallet.status = status.to_string();
|
||||
wallet.total_recharged = total_recharged;
|
||||
wallet.total_consumed = total_consumed;
|
||||
wallet.total_refunded = total_refunded;
|
||||
wallet.total_adjusted = total_adjusted;
|
||||
if let Some(updated_at_unix_secs) = updated_at_unix_secs {
|
||||
wallet.updated_at_unix_secs = updated_at_unix_secs;
|
||||
}
|
||||
return Ok(Some(wallet.clone()));
|
||||
}
|
||||
|
||||
self.data
|
||||
.update_auth_user_wallet_snapshot(
|
||||
user_id,
|
||||
balance,
|
||||
gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
total_recharged,
|
||||
total_consumed,
|
||||
total_refunded,
|
||||
total_adjusted,
|
||||
updated_at_unix_secs,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn update_auth_api_key_wallet_snapshot(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
balance: f64,
|
||||
gift_balance: f64,
|
||||
limit_mode: &str,
|
||||
currency: &str,
|
||||
status: &str,
|
||||
total_recharged: f64,
|
||||
total_consumed: f64,
|
||||
total_refunded: f64,
|
||||
total_adjusted: f64,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Result<Option<aether_data::repository::wallet::StoredWalletSnapshot>, GatewayError> {
|
||||
#[cfg(test)]
|
||||
if let Some(store) = self.auth_wallet_store.as_ref() {
|
||||
let mut guard = store.lock().expect("auth wallet store should lock");
|
||||
let Some((_, wallet)) = guard
|
||||
.iter_mut()
|
||||
.find(|(_, wallet)| wallet.api_key_id.as_deref() == Some(api_key_id))
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
wallet.balance = balance;
|
||||
wallet.gift_balance = gift_balance;
|
||||
wallet.limit_mode = limit_mode.to_string();
|
||||
wallet.currency = currency.to_string();
|
||||
wallet.status = status.to_string();
|
||||
wallet.total_recharged = total_recharged;
|
||||
wallet.total_consumed = total_consumed;
|
||||
wallet.total_refunded = total_refunded;
|
||||
wallet.total_adjusted = total_adjusted;
|
||||
if let Some(updated_at_unix_secs) = updated_at_unix_secs {
|
||||
wallet.updated_at_unix_secs = updated_at_unix_secs;
|
||||
}
|
||||
return Ok(Some(wallet.clone()));
|
||||
}
|
||||
|
||||
self.data
|
||||
.update_auth_api_key_wallet_snapshot(
|
||||
api_key_id,
|
||||
balance,
|
||||
gift_balance,
|
||||
limit_mode,
|
||||
currency,
|
||||
status,
|
||||
total_recharged,
|
||||
total_consumed,
|
||||
total_refunded,
|
||||
total_adjusted,
|
||||
updated_at_unix_secs,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -53,6 +53,15 @@ impl AppState {
|
||||
self.data.has_user_reader()
|
||||
}
|
||||
|
||||
pub fn has_auth_user_data_reader(&self) -> bool {
|
||||
#[cfg(test)]
|
||||
if self.auth_user_store.is_some() {
|
||||
return true;
|
||||
}
|
||||
|
||||
self.data.has_user_reader()
|
||||
}
|
||||
|
||||
pub fn has_usage_data_writer(&self) -> bool {
|
||||
self.data.has_usage_writer()
|
||||
}
|
||||
|
||||
@@ -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("[email protected]".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("[email protected]".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("[email protected]".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": "[email protected]",
|
||||
"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("[email protected]")
|
||||
.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("[email protected]".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("[email protected]".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("[email protected]".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("[email protected]".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