mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat: 扩展 cache creation token 细分统计与 effective_input_tokens 计费逻辑
- 新增 cache_creation_ephemeral_5m/1h_input_tokens 字段,区分不同 TTL 的缓存写入 token - 引入 effective_input_tokens(扣除 cache read 后的有效输入 token),暴露给 usage 接口 - billing 规则生成器支持 5m/1h ephemeral cache 独立定价与分级计费 - usage_mapper 增加 Claude/Anthropic 格式映射,修复 OpenAI responses 格式字段兼容性 - 迁移逻辑增强:支持 checksum 容错、applied/pending 数量日志、逐步执行信息输出 - executor 抽离 LocalExecutionRequestOutcome 类型,统一 sync/stream 路径返回语义 - provider-transport auth 层新增 complete passthrough headers 构建逻辑 - 前端 usage 类型全面补充 effective_input_tokens、cache_creation_tokens、total_input_context 字段
This commit is contained in:
@@ -7,7 +7,9 @@ use crate::ai_pipeline::transport::antigravity::{
|
||||
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
|
||||
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
|
||||
};
|
||||
use crate::ai_pipeline::transport::auth::build_openai_passthrough_headers;
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_complete_passthrough_headers, build_complete_passthrough_headers_with_auth,
|
||||
};
|
||||
use crate::ai_pipeline::transport::claude_code::build_claude_code_passthrough_headers;
|
||||
use crate::ai_pipeline::transport::kiro::{
|
||||
build_kiro_provider_headers, KiroProviderHeadersInput, KIRO_ENVELOPE_NAME,
|
||||
@@ -198,13 +200,13 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
transport.key.fingerprint.as_ref(),
|
||||
)
|
||||
} else if is_vertex {
|
||||
crate::ai_pipeline::transport::build_passthrough_headers(
|
||||
build_complete_passthrough_headers(
|
||||
&parts.headers,
|
||||
&extra_headers,
|
||||
Some("application/json"),
|
||||
)
|
||||
} else {
|
||||
build_openai_passthrough_headers(
|
||||
build_complete_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
auth_header.as_deref().unwrap_or_default(),
|
||||
auth_value.as_deref().unwrap_or_default(),
|
||||
|
||||
@@ -6,7 +6,7 @@ use tracing::warn;
|
||||
use crate::ai_pipeline::planner::common::force_upstream_streaming_for_provider;
|
||||
use crate::ai_pipeline::planner::standard::apply_codex_openai_cli_special_headers;
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_openai_passthrough_headers, ensure_upstream_auth_header,
|
||||
build_claude_passthrough_headers, build_openai_passthrough_headers, ensure_upstream_auth_header,
|
||||
};
|
||||
use crate::ai_pipeline::transport::{
|
||||
apply_local_header_rules, resolve_transport_execution_timeouts,
|
||||
@@ -245,13 +245,23 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
}
|
||||
};
|
||||
|
||||
let mut provider_request_headers = build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
Some("application/json"),
|
||||
);
|
||||
let mut provider_request_headers = if provider_api_format.starts_with("claude:") {
|
||||
build_claude_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
Some("application/json"),
|
||||
)
|
||||
} else {
|
||||
build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
Some("application/json"),
|
||||
)
|
||||
};
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
|
||||
@@ -1,6 +1,15 @@
|
||||
use serde_json::json;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::build_cross_format_openai_cli_request_body;
|
||||
use super::{build_cross_format_openai_cli_request_body, build_local_openai_cli_request_body};
|
||||
|
||||
fn object_keys(value: &Value) -> Vec<&str> {
|
||||
value
|
||||
.as_object()
|
||||
.expect("json object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_openai_chat_cross_format_request_body_from_openai_cli_source() {
|
||||
@@ -26,6 +35,52 @@ fn builds_openai_chat_cross_format_request_body_from_openai_cli_source() {
|
||||
assert_eq!(provider_request_body["messages"][0]["content"], "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_openai_cli_wrapper_preserves_body_order_after_edits() {
|
||||
let body_json: Value = serde_json::from_str(
|
||||
r#"{
|
||||
"text": {"format": {"type": "text"}},
|
||||
"input": [],
|
||||
"model": "gpt-5.4",
|
||||
"store": false,
|
||||
"tools": [],
|
||||
"stream": true,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"reasoning": {"effort": "high"},
|
||||
"tool_choice": "auto"
|
||||
}"#,
|
||||
)
|
||||
.expect("request body should parse");
|
||||
|
||||
let provider_request_body = build_local_openai_cli_request_body(
|
||||
&body_json,
|
||||
"gpt-5.4",
|
||||
true,
|
||||
"codex",
|
||||
"openai:cli",
|
||||
None,
|
||||
Some("key-123"),
|
||||
)
|
||||
.expect("local openai cli body should build");
|
||||
|
||||
assert_eq!(
|
||||
object_keys(&provider_request_body),
|
||||
vec![
|
||||
"text",
|
||||
"input",
|
||||
"model",
|
||||
"store",
|
||||
"tools",
|
||||
"stream",
|
||||
"include",
|
||||
"reasoning",
|
||||
"tool_choice",
|
||||
"instructions",
|
||||
"prompt_cache_key",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_metadata_for_codex_openai_cli_requests() {
|
||||
let body_json = json!({
|
||||
|
||||
@@ -16,7 +16,7 @@ use crate::ai_pipeline::planner::standard::{
|
||||
build_cross_format_openai_chat_upstream_url,
|
||||
};
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_openai_passthrough_headers, ensure_upstream_auth_header,
|
||||
build_claude_passthrough_headers, build_openai_passthrough_headers, ensure_upstream_auth_header,
|
||||
};
|
||||
use crate::ai_pipeline::transport::{
|
||||
apply_local_header_rules, resolve_transport_execution_timeouts,
|
||||
@@ -182,13 +182,23 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut provider_request_headers = build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
Some("application/json"),
|
||||
);
|
||||
let mut provider_request_headers = if provider_api_format.starts_with("claude:") {
|
||||
build_claude_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
Some("application/json"),
|
||||
)
|
||||
} else {
|
||||
build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&BTreeMap::new(),
|
||||
Some("application/json"),
|
||||
)
|
||||
};
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
|
||||
@@ -9,7 +9,8 @@ use crate::ai_pipeline::planner::standard::{
|
||||
build_local_openai_chat_upstream_url,
|
||||
};
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_openai_passthrough_headers, ensure_upstream_auth_header, resolve_local_openai_chat_auth,
|
||||
build_complete_passthrough_headers_with_auth, ensure_upstream_auth_header,
|
||||
resolve_local_openai_chat_auth,
|
||||
};
|
||||
use crate::ai_pipeline::transport::policy::supports_local_openai_chat_transport;
|
||||
use crate::ai_pipeline::transport::{
|
||||
@@ -140,7 +141,7 @@ pub(super) async fn build_same_format_local_openai_chat_decision_payload_for_can
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut provider_request_headers = build_openai_passthrough_headers(
|
||||
let mut provider_request_headers = build_complete_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
|
||||
@@ -22,6 +22,7 @@ use crate::ai_pipeline::transport::antigravity::{
|
||||
};
|
||||
use crate::ai_pipeline::transport::apply_local_header_rules;
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_claude_passthrough_headers, build_complete_passthrough_headers_with_auth,
|
||||
build_openai_passthrough_headers, ensure_upstream_auth_header, resolve_local_gemini_auth,
|
||||
resolve_local_standard_auth,
|
||||
};
|
||||
@@ -354,16 +355,35 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut provider_request_headers = build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&antigravity_auth
|
||||
.as_ref()
|
||||
.map(build_antigravity_static_identity_headers)
|
||||
.unwrap_or_default(),
|
||||
Some("application/json"),
|
||||
);
|
||||
let extra_headers = antigravity_auth
|
||||
.as_ref()
|
||||
.map(build_antigravity_static_identity_headers)
|
||||
.unwrap_or_default();
|
||||
let mut provider_request_headers = if same_format {
|
||||
build_complete_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&extra_headers,
|
||||
Some("application/json"),
|
||||
)
|
||||
} else if provider_api_format.starts_with("claude:") {
|
||||
build_claude_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&extra_headers,
|
||||
Some("application/json"),
|
||||
)
|
||||
} else {
|
||||
build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&extra_headers,
|
||||
Some("application/json"),
|
||||
)
|
||||
};
|
||||
if !apply_local_header_rules(
|
||||
&mut provider_request_headers,
|
||||
transport.endpoint.header_rules.as_ref(),
|
||||
|
||||
@@ -6,6 +6,7 @@ use super::super::{
|
||||
};
|
||||
use crate::ai_pipeline::provider_adaptation_requires_eventstream_accept;
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_claude_passthrough_headers, build_complete_passthrough_headers_with_auth,
|
||||
build_openai_passthrough_headers, ensure_upstream_auth_header,
|
||||
};
|
||||
use crate::ai_pipeline::transport::url::{build_openai_chat_url, build_openai_cli_url};
|
||||
@@ -132,13 +133,31 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
|
||||
};
|
||||
|
||||
let mut provider_request_headers = if payload.provider_request_headers.is_empty() {
|
||||
build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&payload.extra_headers,
|
||||
payload.content_type.as_deref(),
|
||||
)
|
||||
if provider_api_format == client_api_format {
|
||||
build_complete_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&payload.extra_headers,
|
||||
payload.content_type.as_deref(),
|
||||
)
|
||||
} else if provider_api_format.starts_with("claude:") {
|
||||
build_claude_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&payload.extra_headers,
|
||||
payload.content_type.as_deref(),
|
||||
)
|
||||
} else {
|
||||
build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&payload.extra_headers,
|
||||
payload.content_type.as_deref(),
|
||||
)
|
||||
}
|
||||
} else {
|
||||
payload.provider_request_headers.clone()
|
||||
};
|
||||
@@ -281,6 +300,11 @@ pub(crate) fn build_openai_cli_stream_plan_from_decision(
|
||||
} else {
|
||||
provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string());
|
||||
}
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context,
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let plan = ExecutionPlan {
|
||||
request_id,
|
||||
candidate_id: payload.candidate_id.clone(),
|
||||
@@ -306,15 +330,276 @@ pub(crate) fn build_openai_cli_stream_plan_from_decision(
|
||||
timeouts: payload.timeouts.clone(),
|
||||
};
|
||||
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context,
|
||||
&plan.headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
|
||||
Ok(Some(LocalStreamPlanAndReport {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{
|
||||
build_openai_chat_stream_plan_from_decision, build_openai_cli_stream_plan_from_decision,
|
||||
};
|
||||
use crate::GatewayControlSyncDecisionResponse;
|
||||
|
||||
fn object_keys(value: &Value) -> Vec<&str> {
|
||||
value
|
||||
.as_object()
|
||||
.expect("value should be an object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sample_cli_payload() -> GatewayControlSyncDecisionResponse {
|
||||
GatewayControlSyncDecisionResponse {
|
||||
action: "stream".to_string(),
|
||||
decision_kind: Some("openai_cli_stream".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("req_123".to_string()),
|
||||
candidate_id: Some("cand_123".to_string()),
|
||||
provider_name: Some("Codex".to_string()),
|
||||
provider_id: Some("prov_123".to_string()),
|
||||
endpoint_id: Some("ep_123".to_string()),
|
||||
key_id: Some("key_123".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/responses".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer test".to_string()),
|
||||
provider_api_format: Some("openai:cli".to_string()),
|
||||
client_api_format: Some("openai:cli".to_string()),
|
||||
provider_contract: Some("openai:cli".to_string()),
|
||||
client_contract: Some("openai:cli".to_string()),
|
||||
model_name: Some("gpt-5.4".to_string()),
|
||||
mapped_model: Some("gpt-5.4".to_string()),
|
||||
prompt_cache_key: Some("cache-key".to_string()),
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::from([(
|
||||
"content-type".to_string(),
|
||||
"application/json".to_string(),
|
||||
)]),
|
||||
provider_request_body: Some(json!({
|
||||
"text": {"verbosity": "low"},
|
||||
"input": [],
|
||||
"model": "gpt-5.4",
|
||||
"store": false,
|
||||
"tools": [],
|
||||
"stream": true,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"reasoning": {"effort": "high"},
|
||||
"tool_choice": "auto",
|
||||
"instructions": "You are Codex.",
|
||||
"prompt_cache_key": "cache-key"
|
||||
})),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: true,
|
||||
report_kind: Some("openai_cli_stream_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_cli_stream_plan_preserves_provider_request_body_order_in_plan_and_report() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/responses")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = sample_cli_payload();
|
||||
|
||||
let built = build_openai_cli_stream_plan_from_decision(&parts, &json!({}), payload, false)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
let plan_body = built
|
||||
.plan
|
||||
.body
|
||||
.json_body
|
||||
.as_ref()
|
||||
.expect("plan json body should exist");
|
||||
assert_eq!(
|
||||
object_keys(plan_body),
|
||||
vec![
|
||||
"text",
|
||||
"input",
|
||||
"model",
|
||||
"store",
|
||||
"tools",
|
||||
"stream",
|
||||
"include",
|
||||
"reasoning",
|
||||
"tool_choice",
|
||||
"instructions",
|
||||
"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));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_chat_stream_plan_fallback_preserves_complete_same_format_headers() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/chat/completions")
|
||||
.header(http::header::AUTHORIZATION, "Bearer client-token")
|
||||
.header("x-stainless-runtime-version", "v24.0.0")
|
||||
.header("x-app", "codex")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = GatewayControlSyncDecisionResponse {
|
||||
action: "stream".to_string(),
|
||||
decision_kind: Some("openai_chat_stream".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("req_stream_456".to_string()),
|
||||
candidate_id: Some("cand_stream_456".to_string()),
|
||||
provider_name: Some("OpenAI".to_string()),
|
||||
provider_id: Some("prov_stream_456".to_string()),
|
||||
endpoint_id: Some("ep_stream_456".to_string()),
|
||||
key_id: Some("key_stream_456".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/chat/completions".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer upstream-token".to_string()),
|
||||
provider_api_format: Some("openai:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("openai:chat".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some("gpt-5.4".to_string()),
|
||||
mapped_model: Some("gpt-5.4".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::new(),
|
||||
provider_request_body: Some(json!({"model":"gpt-5.4","messages":[],"stream":true})),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: true,
|
||||
report_kind: Some("openai_chat_stream_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
};
|
||||
|
||||
let built = build_openai_chat_stream_plan_from_decision(&parts, &json!({}), payload)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
|
||||
assert_eq!(
|
||||
built.plan.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer upstream-token")
|
||||
);
|
||||
assert_eq!(
|
||||
built
|
||||
.plan
|
||||
.headers
|
||||
.get("x-stainless-runtime-version")
|
||||
.map(String::as_str),
|
||||
Some("v24.0.0")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("x-app").map(String::as_str),
|
||||
Some("codex")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("accept").map(String::as_str),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_chat_stream_plan_fallback_restores_claude_headers_for_cross_format() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/chat/completions")
|
||||
.header("anthropic-beta", "prompt-caching-2024-07-31")
|
||||
.header("x-stainless-runtime-version", "v24.0.0")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = GatewayControlSyncDecisionResponse {
|
||||
action: "stream".to_string(),
|
||||
decision_kind: Some("openai_chat_stream".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: Some("format_conversion".to_string()),
|
||||
request_id: Some("req_stream_789".to_string()),
|
||||
candidate_id: Some("cand_stream_789".to_string()),
|
||||
provider_name: Some("Claude".to_string()),
|
||||
provider_id: Some("prov_stream_789".to_string()),
|
||||
endpoint_id: Some("ep_stream_789".to_string()),
|
||||
key_id: Some("key_stream_789".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/messages".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("x-api-key".to_string()),
|
||||
auth_value: Some("sk-upstream-claude".to_string()),
|
||||
provider_api_format: Some("claude:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("claude:chat".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some("claude-sonnet-4-5".to_string()),
|
||||
mapped_model: Some("claude-sonnet-4-5".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::new(),
|
||||
provider_request_body: Some(
|
||||
json!({"model":"claude-sonnet-4-5","messages":[],"stream":true}),
|
||||
),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: true,
|
||||
report_kind: Some("openai_chat_stream_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
};
|
||||
|
||||
let built = build_openai_chat_stream_plan_from_decision(&parts, &json!({}), payload)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
|
||||
assert_eq!(
|
||||
built.plan.headers.get("x-api-key").map(String::as_str),
|
||||
Some("sk-upstream-claude")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("anthropic-beta").map(String::as_str),
|
||||
Some("prompt-caching-2024-07-31")
|
||||
);
|
||||
assert_eq!(
|
||||
built
|
||||
.plan
|
||||
.headers
|
||||
.get("anthropic-version")
|
||||
.map(String::as_str),
|
||||
Some("2023-06-01")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("accept").map(String::as_str),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ use super::super::{
|
||||
LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::ai_pipeline::transport::auth::{
|
||||
build_claude_passthrough_headers, build_complete_passthrough_headers_with_auth,
|
||||
build_openai_passthrough_headers, ensure_upstream_auth_header,
|
||||
};
|
||||
use crate::ai_pipeline::transport::url::{build_openai_chat_url, build_openai_cli_url};
|
||||
@@ -132,13 +133,31 @@ pub(crate) fn build_openai_chat_sync_plan_from_decision(
|
||||
};
|
||||
|
||||
let mut provider_request_headers = if payload.provider_request_headers.is_empty() {
|
||||
build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&payload.extra_headers,
|
||||
payload.content_type.as_deref(),
|
||||
)
|
||||
if provider_api_format == client_api_format {
|
||||
build_complete_passthrough_headers_with_auth(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&payload.extra_headers,
|
||||
payload.content_type.as_deref(),
|
||||
)
|
||||
} else if provider_api_format.starts_with("claude:") {
|
||||
build_claude_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&payload.extra_headers,
|
||||
payload.content_type.as_deref(),
|
||||
)
|
||||
} else {
|
||||
build_openai_passthrough_headers(
|
||||
&parts.headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
&payload.extra_headers,
|
||||
payload.content_type.as_deref(),
|
||||
)
|
||||
}
|
||||
} else {
|
||||
payload.provider_request_headers.clone()
|
||||
};
|
||||
@@ -275,6 +294,11 @@ pub(crate) fn build_openai_cli_sync_plan_from_decision(
|
||||
if payload.upstream_is_stream && !provider_request_headers.contains_key("accept") {
|
||||
provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string());
|
||||
}
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context,
|
||||
&provider_request_headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
let plan = ExecutionPlan {
|
||||
request_id,
|
||||
candidate_id: payload.candidate_id.clone(),
|
||||
@@ -300,15 +324,268 @@ pub(crate) fn build_openai_cli_sync_plan_from_decision(
|
||||
timeouts: payload.timeouts.clone(),
|
||||
};
|
||||
|
||||
let report_context = augment_sync_report_context(
|
||||
payload.report_context,
|
||||
&plan.headers,
|
||||
&provider_request_body_value,
|
||||
)?;
|
||||
|
||||
Ok(Some(LocalSyncPlanAndReport {
|
||||
plan,
|
||||
report_kind: payload.report_kind,
|
||||
report_context,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{
|
||||
build_openai_chat_sync_plan_from_decision, build_openai_cli_sync_plan_from_decision,
|
||||
};
|
||||
use crate::GatewayControlSyncDecisionResponse;
|
||||
|
||||
fn object_keys(value: &Value) -> Vec<&str> {
|
||||
value
|
||||
.as_object()
|
||||
.expect("value should be an object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sample_cli_payload() -> GatewayControlSyncDecisionResponse {
|
||||
GatewayControlSyncDecisionResponse {
|
||||
action: "sync".to_string(),
|
||||
decision_kind: Some("openai_cli_sync".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("req_123".to_string()),
|
||||
candidate_id: Some("cand_123".to_string()),
|
||||
provider_name: Some("Codex".to_string()),
|
||||
provider_id: Some("prov_123".to_string()),
|
||||
endpoint_id: Some("ep_123".to_string()),
|
||||
key_id: Some("key_123".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/responses".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer test".to_string()),
|
||||
provider_api_format: Some("openai:cli".to_string()),
|
||||
client_api_format: Some("openai:cli".to_string()),
|
||||
provider_contract: Some("openai:cli".to_string()),
|
||||
client_contract: Some("openai:cli".to_string()),
|
||||
model_name: Some("gpt-5.4".to_string()),
|
||||
mapped_model: Some("gpt-5.4".to_string()),
|
||||
prompt_cache_key: Some("cache-key".to_string()),
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::from([(
|
||||
"content-type".to_string(),
|
||||
"application/json".to_string(),
|
||||
)]),
|
||||
provider_request_body: Some(json!({
|
||||
"text": {"verbosity": "low"},
|
||||
"input": [],
|
||||
"model": "gpt-5.4",
|
||||
"store": false,
|
||||
"tools": [],
|
||||
"stream": true,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"reasoning": {"effort": "high"},
|
||||
"tool_choice": "auto",
|
||||
"instructions": "You are Codex.",
|
||||
"prompt_cache_key": "cache-key"
|
||||
})),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: true,
|
||||
report_kind: Some("openai_cli_sync_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_cli_sync_plan_preserves_provider_request_body_order_in_plan_and_report() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/responses")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = sample_cli_payload();
|
||||
|
||||
let built = build_openai_cli_sync_plan_from_decision(&parts, &json!({}), payload, false)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
let plan_body = built
|
||||
.plan
|
||||
.body
|
||||
.json_body
|
||||
.as_ref()
|
||||
.expect("plan json body should exist");
|
||||
assert_eq!(
|
||||
object_keys(plan_body),
|
||||
vec![
|
||||
"text",
|
||||
"input",
|
||||
"model",
|
||||
"store",
|
||||
"tools",
|
||||
"stream",
|
||||
"include",
|
||||
"reasoning",
|
||||
"tool_choice",
|
||||
"instructions",
|
||||
"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));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_chat_sync_plan_fallback_preserves_complete_same_format_headers() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/chat/completions")
|
||||
.header(http::header::AUTHORIZATION, "Bearer client-token")
|
||||
.header("x-stainless-runtime-version", "v24.0.0")
|
||||
.header("x-app", "codex")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = GatewayControlSyncDecisionResponse {
|
||||
action: "sync".to_string(),
|
||||
decision_kind: Some("openai_chat_sync".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("req_456".to_string()),
|
||||
candidate_id: Some("cand_456".to_string()),
|
||||
provider_name: Some("OpenAI".to_string()),
|
||||
provider_id: Some("prov_456".to_string()),
|
||||
endpoint_id: Some("ep_456".to_string()),
|
||||
key_id: Some("key_456".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/chat/completions".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("authorization".to_string()),
|
||||
auth_value: Some("Bearer upstream-token".to_string()),
|
||||
provider_api_format: Some("openai:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("openai:chat".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some("gpt-5.4".to_string()),
|
||||
mapped_model: Some("gpt-5.4".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::new(),
|
||||
provider_request_body: Some(json!({"model":"gpt-5.4","messages":[],"stream":false})),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: false,
|
||||
report_kind: Some("openai_chat_sync_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
};
|
||||
|
||||
let built = build_openai_chat_sync_plan_from_decision(&parts, &json!({}), payload)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
|
||||
assert_eq!(
|
||||
built.plan.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer upstream-token")
|
||||
);
|
||||
assert_eq!(
|
||||
built
|
||||
.plan
|
||||
.headers
|
||||
.get("x-stainless-runtime-version")
|
||||
.map(String::as_str),
|
||||
Some("v24.0.0")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("x-app").map(String::as_str),
|
||||
Some("codex")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_openai_chat_sync_plan_fallback_restores_claude_headers_for_cross_format() {
|
||||
let parts = http::Request::builder()
|
||||
.uri("http://localhost/v1/chat/completions")
|
||||
.header("anthropic-beta", "prompt-caching-2024-07-31")
|
||||
.header("x-stainless-runtime-version", "v24.0.0")
|
||||
.body(())
|
||||
.expect("request should build")
|
||||
.into_parts()
|
||||
.0;
|
||||
let payload = GatewayControlSyncDecisionResponse {
|
||||
action: "sync".to_string(),
|
||||
decision_kind: Some("openai_chat_sync".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: Some("format_conversion".to_string()),
|
||||
request_id: Some("req_789".to_string()),
|
||||
candidate_id: Some("cand_789".to_string()),
|
||||
provider_name: Some("Claude".to_string()),
|
||||
provider_id: Some("prov_789".to_string()),
|
||||
endpoint_id: Some("ep_789".to_string()),
|
||||
key_id: Some("key_789".to_string()),
|
||||
upstream_base_url: Some("https://example.com".to_string()),
|
||||
upstream_url: Some("https://example.com/v1/messages".to_string()),
|
||||
provider_request_method: None,
|
||||
auth_header: Some("x-api-key".to_string()),
|
||||
auth_value: Some("sk-upstream-claude".to_string()),
|
||||
provider_api_format: Some("claude:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: Some("claude:chat".to_string()),
|
||||
client_contract: Some("openai:chat".to_string()),
|
||||
model_name: Some("claude-sonnet-4-5".to_string()),
|
||||
mapped_model: Some("claude-sonnet-4-5".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::new(),
|
||||
provider_request_body: Some(
|
||||
json!({"model":"claude-sonnet-4-5","messages":[],"stream":false}),
|
||||
),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: false,
|
||||
report_kind: Some("openai_chat_sync_success".to_string()),
|
||||
report_context: Some(json!({})),
|
||||
auth_context: None,
|
||||
};
|
||||
|
||||
let built = build_openai_chat_sync_plan_from_decision(&parts, &json!({}), payload)
|
||||
.expect("plan build should succeed")
|
||||
.expect("plan should be produced");
|
||||
|
||||
assert_eq!(
|
||||
built.plan.headers.get("x-api-key").map(String::as_str),
|
||||
Some("sk-upstream-claude")
|
||||
);
|
||||
assert_eq!(
|
||||
built.plan.headers.get("anthropic-beta").map(String::as_str),
|
||||
Some("prompt-caching-2024-07-31")
|
||||
);
|
||||
assert_eq!(
|
||||
built
|
||||
.plan
|
||||
.headers
|
||||
.get("anthropic-version")
|
||||
.map(String::as_str),
|
||||
Some("2023-06-01")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,9 @@ use axum::http::{HeaderName, HeaderValue, Response};
|
||||
|
||||
use crate::constants::CONTROL_EXECUTED_HEADER;
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::executor::{maybe_execute_stream_request, maybe_execute_sync_request};
|
||||
use crate::executor::{
|
||||
maybe_execute_stream_request, maybe_execute_sync_request, LocalExecutionRequestOutcome,
|
||||
};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::resolve_execution_runtime_auth_context;
|
||||
@@ -19,9 +21,9 @@ pub(crate) async fn maybe_execute_via_control(
|
||||
trace_id: &str,
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
require_stream: bool,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some(decision) = decision else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
let mut local_decision = decision.clone();
|
||||
@@ -46,7 +48,15 @@ pub(crate) async fn maybe_execute_via_control(
|
||||
.await?
|
||||
};
|
||||
|
||||
Ok(response.map(mark_control_executed))
|
||||
Ok(match response {
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
LocalExecutionRequestOutcome::Responded(mark_control_executed(response))
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => {
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome)
|
||||
}
|
||||
LocalExecutionRequestOutcome::NoPath => LocalExecutionRequestOutcome::NoPath,
|
||||
})
|
||||
}
|
||||
|
||||
fn mark_control_executed(mut response: Response<Body>) -> Response<Body> {
|
||||
|
||||
@@ -1,12 +1,10 @@
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
|
||||
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
|
||||
use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate;
|
||||
|
||||
use crate::ai_pipeline_api::{LocalStreamPlanAndReport, LocalSyncPlanAndReport};
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::execution_runtime::{execute_execution_runtime_stream, execute_execution_runtime_sync};
|
||||
use crate::executor::{build_local_execution_exhaustion, LocalExecutionRequestOutcome};
|
||||
use crate::request_candidate_runtime::record_local_request_candidate_status;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
@@ -53,12 +51,17 @@ pub(crate) async fn execute_sync_plan_and_reports<T>(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
plan_and_reports: Vec<T>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError>
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError>
|
||||
where
|
||||
T: LocalPlanAndReport,
|
||||
{
|
||||
let mut remaining = plan_and_reports.into_iter();
|
||||
let mut last_attempted = None;
|
||||
while let Some(plan_and_report) = remaining.next() {
|
||||
last_attempted = Some((
|
||||
plan_and_report.plan().clone(),
|
||||
plan_and_report.report_context(),
|
||||
));
|
||||
if let Some(response) = execute_execution_runtime_sync(
|
||||
state,
|
||||
parts.uri.path(),
|
||||
@@ -72,11 +75,16 @@ where
|
||||
.await?
|
||||
{
|
||||
mark_unused_local_candidates(state, remaining.collect()).await;
|
||||
return Ok(Some(response));
|
||||
return Ok(LocalExecutionRequestOutcome::responded(response));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
let Some((plan, report_context)) = last_attempted else {
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
Ok(LocalExecutionRequestOutcome::Exhausted(
|
||||
build_local_execution_exhaustion(state, &plan, report_context.as_ref()).await,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn execute_stream_plan_and_reports<T>(
|
||||
@@ -85,12 +93,17 @@ pub(crate) async fn execute_stream_plan_and_reports<T>(
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
plan_and_reports: Vec<T>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError>
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError>
|
||||
where
|
||||
T: LocalPlanAndReport,
|
||||
{
|
||||
let mut remaining = plan_and_reports.into_iter();
|
||||
let mut last_attempted = None;
|
||||
while let Some(plan_and_report) = remaining.next() {
|
||||
last_attempted = Some((
|
||||
plan_and_report.plan().clone(),
|
||||
plan_and_report.report_context(),
|
||||
));
|
||||
if let Some(response) = execute_execution_runtime_stream(
|
||||
state,
|
||||
plan_and_report.plan().clone(),
|
||||
@@ -103,11 +116,16 @@ where
|
||||
.await?
|
||||
{
|
||||
mark_unused_local_candidates(state, remaining.collect()).await;
|
||||
return Ok(Some(response));
|
||||
return Ok(LocalExecutionRequestOutcome::responded(response));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
let Some((plan, report_context)) = last_attempted else {
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
Ok(LocalExecutionRequestOutcome::Exhausted(
|
||||
build_local_execution_exhaustion(state, &plan, report_context.as_ref()).await,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_unused_local_candidates<T>(state: &AppState, remaining: Vec<T>)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
pub(crate) mod candidate_loop;
|
||||
mod orchestration;
|
||||
mod outcome;
|
||||
mod plan_fallback;
|
||||
mod policy;
|
||||
mod remote;
|
||||
@@ -9,8 +10,15 @@ mod sync_path;
|
||||
pub(crate) use crate::request_candidate_runtime::{
|
||||
persist_available_local_candidate, persist_skipped_local_candidate,
|
||||
};
|
||||
pub(crate) use candidate_loop::mark_unused_local_candidate_items;
|
||||
pub(crate) use candidate_loop::{
|
||||
execute_stream_plan_and_reports, execute_sync_plan_and_reports,
|
||||
mark_unused_local_candidate_items,
|
||||
};
|
||||
pub(crate) use orchestration::*;
|
||||
pub(crate) use outcome::{
|
||||
build_local_execution_exhaustion, record_failed_usage_for_exhausted_request,
|
||||
LocalExecutionExhaustion, LocalExecutionRequestOutcome,
|
||||
};
|
||||
pub(crate) use plan_fallback::{
|
||||
maybe_execute_stream_via_plan_fallback, maybe_execute_sync_via_plan_fallback,
|
||||
};
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
|
||||
use crate::ai_pipeline_api::{
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind,
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind,
|
||||
@@ -21,6 +18,7 @@ use crate::control::GatewayControlDecision;
|
||||
use crate::executor::candidate_loop::{
|
||||
execute_stream_plan_and_reports, execute_sync_plan_and_reports,
|
||||
};
|
||||
use crate::executor::LocalExecutionRequestOutcome;
|
||||
use crate::{AppState, GatewayControlSyncDecisionResponse, GatewayError};
|
||||
|
||||
pub(crate) async fn maybe_execute_sync_local_path(
|
||||
@@ -29,7 +27,7 @@ pub(crate) async fn maybe_execute_sync_local_path(
|
||||
body_bytes: &axum::body::Bytes,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
super::maybe_execute_via_sync_decision_path(state, parts, body_bytes, trace_id, decision).await
|
||||
}
|
||||
|
||||
@@ -39,7 +37,7 @@ pub(crate) async fn maybe_execute_stream_local_path(
|
||||
body_bytes: &axum::body::Bytes,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
super::maybe_execute_via_stream_decision_path(state, parts, body_bytes, trace_id, decision)
|
||||
.await
|
||||
}
|
||||
@@ -51,17 +49,17 @@ pub(crate) async fn maybe_execute_sync_via_local_decision(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let plan_and_reports = build_local_openai_chat_sync_plan_and_reports_for_kind(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
let plan_count = plan_and_reports.len();
|
||||
if let Some(response) = execute_sync_plan_and_reports(
|
||||
let outcome = execute_sync_plan_and_reports(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
@@ -69,15 +67,15 @@ pub(crate) async fn maybe_execute_sync_via_local_decision(
|
||||
plan_kind,
|
||||
plan_and_reports,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
.await?;
|
||||
|
||||
if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome {
|
||||
set_local_openai_chat_execution_exhausted_diagnostic(
|
||||
state, trace_id, decision, plan_kind, body_json, plan_count,
|
||||
);
|
||||
}
|
||||
|
||||
set_local_openai_chat_execution_exhausted_diagnostic(
|
||||
state, trace_id, decision, plan_kind, body_json, plan_count,
|
||||
);
|
||||
Ok(None)
|
||||
Ok(outcome)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_execute_stream_via_local_decision(
|
||||
@@ -87,27 +85,27 @@ pub(crate) async fn maybe_execute_stream_via_local_decision(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let plan_and_reports = build_local_openai_chat_stream_plan_and_reports_for_kind(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
let plan_count = plan_and_reports.len();
|
||||
if let Some(response) =
|
||||
let outcome =
|
||||
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
.await?;
|
||||
|
||||
if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome {
|
||||
set_local_openai_chat_execution_exhausted_diagnostic(
|
||||
state, trace_id, decision, plan_kind, body_json, plan_count,
|
||||
);
|
||||
}
|
||||
|
||||
set_local_openai_chat_execution_exhausted_diagnostic(
|
||||
state, trace_id, decision, plan_kind, body_json, plan_count,
|
||||
);
|
||||
Ok(None)
|
||||
Ok(outcome)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_execute_sync_via_local_openai_cli_decision(
|
||||
@@ -117,14 +115,14 @@ pub(crate) async fn maybe_execute_sync_via_local_openai_cli_decision(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
|
||||
build_local_openai_cli_sync_plan_and_reports_for_kind(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
execute_sync_plan_and_reports(
|
||||
@@ -145,14 +143,14 @@ pub(crate) async fn maybe_execute_stream_via_local_openai_cli_decision(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let plan_and_reports: Vec<LocalStreamPlanAndReport> =
|
||||
build_local_openai_cli_stream_plan_and_reports_for_kind(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
|
||||
@@ -166,9 +164,9 @@ pub(crate) async fn maybe_execute_sync_via_standard_family_decision(
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
resolve_sync_spec: fn(&str) -> Option<LocalStandardSpec>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some(spec) = resolve_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
|
||||
@@ -177,7 +175,7 @@ pub(crate) async fn maybe_execute_sync_via_standard_family_decision(
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
execute_sync_plan_and_reports(
|
||||
@@ -199,9 +197,9 @@ pub(crate) async fn maybe_execute_stream_via_standard_family_decision(
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
resolve_stream_spec: fn(&str) -> Option<LocalStandardSpec>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some(spec) = resolve_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
let plan_and_reports: Vec<LocalStreamPlanAndReport> =
|
||||
@@ -210,7 +208,7 @@ pub(crate) async fn maybe_execute_stream_via_standard_family_decision(
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
|
||||
@@ -223,8 +221,10 @@ pub(crate) async fn maybe_execute_sync_via_local_standard_decision(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
if let Some(response) = maybe_execute_sync_via_standard_family_decision(
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let mut exhausted = None;
|
||||
|
||||
match maybe_execute_sync_via_standard_family_decision(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
@@ -235,10 +235,14 @@ pub(crate) async fn maybe_execute_sync_via_local_standard_decision(
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
maybe_execute_sync_via_standard_family_decision(
|
||||
match maybe_execute_sync_via_standard_family_decision(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
@@ -247,7 +251,18 @@ pub(crate) async fn maybe_execute_sync_via_local_standard_decision(
|
||||
plan_kind,
|
||||
resolve_gemini_sync_spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
Ok(LocalExecutionRequestOutcome::Responded(response))
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => {
|
||||
Ok(LocalExecutionRequestOutcome::Exhausted(outcome))
|
||||
}
|
||||
LocalExecutionRequestOutcome::NoPath => Ok(exhausted
|
||||
.map(LocalExecutionRequestOutcome::Exhausted)
|
||||
.unwrap_or(LocalExecutionRequestOutcome::NoPath)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_execute_stream_via_local_standard_decision(
|
||||
@@ -257,8 +272,10 @@ pub(crate) async fn maybe_execute_stream_via_local_standard_decision(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
if let Some(response) = maybe_execute_stream_via_standard_family_decision(
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let mut exhausted = None;
|
||||
|
||||
match maybe_execute_stream_via_standard_family_decision(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
@@ -269,10 +286,14 @@ pub(crate) async fn maybe_execute_stream_via_local_standard_decision(
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
maybe_execute_stream_via_standard_family_decision(
|
||||
match maybe_execute_stream_via_standard_family_decision(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
@@ -281,7 +302,18 @@ pub(crate) async fn maybe_execute_stream_via_local_standard_decision(
|
||||
plan_kind,
|
||||
resolve_gemini_stream_spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
Ok(LocalExecutionRequestOutcome::Responded(response))
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => {
|
||||
Ok(LocalExecutionRequestOutcome::Exhausted(outcome))
|
||||
}
|
||||
LocalExecutionRequestOutcome::NoPath => Ok(exhausted
|
||||
.map(LocalExecutionRequestOutcome::Exhausted)
|
||||
.unwrap_or(LocalExecutionRequestOutcome::NoPath)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_execute_sync_via_local_same_format_provider_decision(
|
||||
@@ -291,9 +323,9 @@ pub(crate) async fn maybe_execute_sync_via_local_same_format_provider_decision(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some(spec) = resolve_local_same_format_sync_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
|
||||
@@ -302,7 +334,7 @@ pub(crate) async fn maybe_execute_sync_via_local_same_format_provider_decision(
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
execute_sync_plan_and_reports(
|
||||
@@ -323,9 +355,9 @@ pub(crate) async fn maybe_execute_stream_via_local_same_format_provider_decision
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some(spec) = resolve_local_same_format_stream_spec(plan_kind) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
let plan_and_reports: Vec<LocalStreamPlanAndReport> =
|
||||
@@ -334,7 +366,7 @@ pub(crate) async fn maybe_execute_stream_via_local_same_format_provider_decision
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
|
||||
@@ -349,7 +381,7 @@ pub(crate) async fn maybe_execute_sync_via_local_gemini_files_decision(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
|
||||
build_local_gemini_files_sync_plan_and_reports_for_kind(
|
||||
state,
|
||||
@@ -363,7 +395,7 @@ pub(crate) async fn maybe_execute_sync_via_local_gemini_files_decision(
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
execute_sync_plan_and_reports(
|
||||
@@ -383,14 +415,14 @@ pub(crate) async fn maybe_execute_stream_via_local_gemini_files_decision(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let plan_and_reports: Vec<LocalStreamPlanAndReport> =
|
||||
build_local_gemini_files_stream_plan_and_reports_for_kind(
|
||||
state, parts, trace_id, decision, plan_kind,
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
|
||||
@@ -403,14 +435,14 @@ pub(crate) async fn maybe_execute_sync_via_local_video_decision(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let plan_and_reports: Vec<LocalSyncPlanAndReport> =
|
||||
build_local_video_sync_plan_and_reports_for_kind(
|
||||
state, parts, body_json, trace_id, decision, plan_kind,
|
||||
)
|
||||
.await?;
|
||||
if plan_and_reports.is_empty() {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
execute_sync_plan_and_reports(
|
||||
@@ -430,14 +462,14 @@ pub(crate) async fn maybe_execute_sync_request(
|
||||
body_bytes: &axum::body::Bytes,
|
||||
trace_id: &str,
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some(decision) = decision else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
#[cfg(not(test))]
|
||||
{
|
||||
if parts.method != http::Method::POST {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
return maybe_execute_sync_local_path(state, parts, body_bytes, trace_id, decision).await;
|
||||
}
|
||||
@@ -449,7 +481,7 @@ pub(crate) async fn maybe_execute_sync_request(
|
||||
.is_empty()
|
||||
&& parts.method != http::Method::POST
|
||||
{
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
maybe_execute_sync_local_path(state, parts, body_bytes, trace_id, decision).await
|
||||
}
|
||||
@@ -461,14 +493,14 @@ pub(crate) async fn maybe_execute_stream_request(
|
||||
body_bytes: &axum::body::Bytes,
|
||||
trace_id: &str,
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some(decision) = decision else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
#[cfg(not(test))]
|
||||
{
|
||||
if parts.method != http::Method::POST {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
return maybe_execute_stream_local_path(state, parts, body_bytes, trace_id, decision).await;
|
||||
}
|
||||
@@ -480,7 +512,7 @@ pub(crate) async fn maybe_execute_stream_request(
|
||||
.is_empty()
|
||||
&& parts.method != http::Method::POST
|
||||
{
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
maybe_execute_stream_local_path(state, parts, body_bytes, trace_id, decision).await
|
||||
}
|
||||
|
||||
232
apps/aether-gateway/src/executor/outcome.rs
Normal file
232
apps/aether-gateway/src/executor/outcome.rs
Normal file
@@ -0,0 +1,232 @@
|
||||
use std::time::Instant;
|
||||
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use aether_usage_runtime::{
|
||||
build_usage_event_data_seed, UsageEvent, UsageEventData, UsageEventType,
|
||||
};
|
||||
use axum::body::Body;
|
||||
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::state::LocalExecutionRuntimeMissDiagnostic;
|
||||
use crate::AppState;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) enum LocalExecutionRequestOutcome {
|
||||
Responded(Response<Body>),
|
||||
Exhausted(LocalExecutionExhaustion),
|
||||
NoPath,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalExecutionExhaustion {
|
||||
request_id: String,
|
||||
data: UsageEventData,
|
||||
candidate_id: Option<String>,
|
||||
candidate_index: Option<u32>,
|
||||
upstream_status_code: Option<u16>,
|
||||
upstream_error_type: Option<String>,
|
||||
upstream_error_message: Option<String>,
|
||||
}
|
||||
|
||||
impl LocalExecutionRequestOutcome {
|
||||
pub(crate) fn responded(response: Response<Body>) -> Self {
|
||||
Self::Responded(response)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn build_local_execution_exhaustion(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
) -> LocalExecutionExhaustion {
|
||||
let mut data = build_usage_event_data_seed(plan, report_context);
|
||||
let last_failed_candidate = match state
|
||||
.read_request_candidates_by_request_id(plan.request_id.as_str())
|
||||
.await
|
||||
{
|
||||
Ok(candidates) => select_last_failed_request_candidate(&candidates).cloned(),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
request_id = %plan.request_id,
|
||||
error = ?err,
|
||||
"gateway failed to load request candidates for exhausted local execution"
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(candidate) = last_failed_candidate.as_ref() {
|
||||
data.user_id = data.user_id.or_else(|| candidate.user_id.clone());
|
||||
data.api_key_id = data.api_key_id.or_else(|| candidate.api_key_id.clone());
|
||||
data.username = data.username.or_else(|| candidate.username.clone());
|
||||
data.api_key_name = data.api_key_name.or_else(|| candidate.api_key_name.clone());
|
||||
data.provider_id = data.provider_id.or_else(|| candidate.provider_id.clone());
|
||||
data.provider_endpoint_id = data
|
||||
.provider_endpoint_id
|
||||
.or_else(|| candidate.endpoint_id.clone());
|
||||
data.provider_api_key_id = data
|
||||
.provider_api_key_id
|
||||
.or_else(|| candidate.key_id.clone());
|
||||
}
|
||||
|
||||
LocalExecutionExhaustion {
|
||||
request_id: plan.request_id.clone(),
|
||||
data,
|
||||
candidate_id: last_failed_candidate
|
||||
.as_ref()
|
||||
.map(|candidate| candidate.id.clone()),
|
||||
candidate_index: last_failed_candidate
|
||||
.as_ref()
|
||||
.map(|candidate| candidate.candidate_index),
|
||||
upstream_status_code: last_failed_candidate
|
||||
.as_ref()
|
||||
.and_then(|candidate| candidate.status_code),
|
||||
upstream_error_type: last_failed_candidate
|
||||
.as_ref()
|
||||
.and_then(|candidate| candidate.error_type.clone())
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty()),
|
||||
upstream_error_message: last_failed_candidate
|
||||
.as_ref()
|
||||
.and_then(|candidate| candidate.error_message.clone())
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn record_failed_usage_for_exhausted_request(
|
||||
state: &AppState,
|
||||
exhaustion: LocalExecutionExhaustion,
|
||||
started_at: &Instant,
|
||||
local_execution_runtime_miss_detail: &str,
|
||||
diagnostic: Option<&LocalExecutionRuntimeMissDiagnostic>,
|
||||
) {
|
||||
if !state.usage_runtime.is_enabled() {
|
||||
return;
|
||||
}
|
||||
|
||||
let LocalExecutionExhaustion {
|
||||
request_id,
|
||||
mut data,
|
||||
candidate_id,
|
||||
candidate_index,
|
||||
upstream_status_code,
|
||||
upstream_error_type,
|
||||
upstream_error_message,
|
||||
} = exhaustion;
|
||||
|
||||
let status_code = http::StatusCode::SERVICE_UNAVAILABLE.as_u16();
|
||||
let candidate_status_code = upstream_status_code.unwrap_or(status_code);
|
||||
data.status_code = Some(status_code);
|
||||
data.error_message = upstream_error_message
|
||||
.clone()
|
||||
.or_else(|| Some(local_execution_runtime_miss_detail.to_string()));
|
||||
data.error_category = error_category_for_failed_status(status_code);
|
||||
data.response_time_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
data.response_headers = Some(json_header_map());
|
||||
data.response_body = Some(json!({
|
||||
"error": {
|
||||
"type": upstream_error_type
|
||||
.as_deref()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.unwrap_or("upstream_error"),
|
||||
"message": upstream_error_message
|
||||
.as_deref()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.unwrap_or(local_execution_runtime_miss_detail),
|
||||
"code": candidate_status_code,
|
||||
}
|
||||
}));
|
||||
|
||||
let mut client_headers = Map::from_iter([(
|
||||
"content-type".to_string(),
|
||||
Value::String("application/json".to_string()),
|
||||
)]);
|
||||
if let Some(reason) = diagnostic
|
||||
.and_then(|value| Some(value.reason.trim()))
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
client_headers.insert(
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER.to_string(),
|
||||
Value::String(reason.to_string()),
|
||||
);
|
||||
}
|
||||
data.client_response_headers = Some(Value::Object(client_headers));
|
||||
data.client_response_body = Some(json!({
|
||||
"error": {
|
||||
"type": "http_error",
|
||||
"message": local_execution_runtime_miss_detail,
|
||||
}
|
||||
}));
|
||||
|
||||
let mut request_metadata = match data.request_metadata.take() {
|
||||
Some(Value::Object(object)) => object,
|
||||
Some(other) => Map::from_iter([("seed".to_string(), other)]),
|
||||
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()),
|
||||
);
|
||||
}
|
||||
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;
|
||||
}
|
||||
|
||||
fn select_last_failed_request_candidate(
|
||||
candidates: &[StoredRequestCandidate],
|
||||
) -> Option<&StoredRequestCandidate> {
|
||||
candidates
|
||||
.iter()
|
||||
.filter(|candidate| {
|
||||
matches!(
|
||||
candidate.status,
|
||||
RequestCandidateStatus::Failed | RequestCandidateStatus::Cancelled
|
||||
)
|
||||
})
|
||||
.max_by_key(|candidate| {
|
||||
(
|
||||
candidate.retry_index,
|
||||
candidate.candidate_index,
|
||||
candidate
|
||||
.finished_at_unix_ms
|
||||
.or(candidate.started_at_unix_ms)
|
||||
.unwrap_or(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())
|
||||
} else if status_code >= 400 {
|
||||
Some("client_error".to_string())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn json_header_map() -> Value {
|
||||
Value::Object(Map::from_iter([(
|
||||
"content-type".to_string(),
|
||||
Value::String("application/json".to_string()),
|
||||
)]))
|
||||
}
|
||||
@@ -1,9 +1,11 @@
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
|
||||
use crate::ai_pipeline_api::{maybe_build_stream_plan_payload, maybe_build_sync_plan_payload};
|
||||
use crate::ai_pipeline_api::{
|
||||
maybe_build_stream_plan_payload, maybe_build_sync_plan_payload, LocalStreamPlanAndReport,
|
||||
LocalSyncPlanAndReport,
|
||||
};
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::execution_runtime::{execute_execution_runtime_stream, execute_execution_runtime_sync};
|
||||
use crate::executor::{
|
||||
execute_stream_plan_and_reports, execute_sync_plan_and_reports, LocalExecutionRequestOutcome,
|
||||
};
|
||||
use crate::{AppState, GatewayControlPlanResponse, GatewayError, GatewayFallbackReason};
|
||||
|
||||
pub(crate) async fn maybe_execute_sync_via_plan_fallback(
|
||||
@@ -16,7 +18,7 @@ pub(crate) async fn maybe_execute_sync_via_plan_fallback(
|
||||
_plan_kind: &str,
|
||||
_bypass_cache_key: String,
|
||||
_fallback_reason: GatewayFallbackReason,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let body_is_empty =
|
||||
body_base64.is_none() && body_json.as_object().is_some_and(|value| value.is_empty());
|
||||
let Some(payload) = maybe_build_sync_plan_payload(
|
||||
@@ -30,7 +32,7 @@ pub(crate) async fn maybe_execute_sync_via_plan_fallback(
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
let GatewayControlPlanResponse {
|
||||
@@ -43,18 +45,20 @@ pub(crate) async fn maybe_execute_sync_via_plan_fallback(
|
||||
} = payload;
|
||||
|
||||
let (Some(plan_kind), Some(plan)) = (plan_kind, plan) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
execute_execution_runtime_sync(
|
||||
execute_sync_plan_and_reports(
|
||||
state,
|
||||
parts.uri.path(),
|
||||
plan,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind.as_str(),
|
||||
report_kind,
|
||||
report_context,
|
||||
vec![LocalSyncPlanAndReport {
|
||||
plan,
|
||||
report_kind,
|
||||
report_context,
|
||||
}],
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -69,11 +73,11 @@ pub(crate) async fn maybe_execute_stream_via_plan_fallback(
|
||||
_plan_kind: &str,
|
||||
_bypass_cache_key: String,
|
||||
_fallback_reason: GatewayFallbackReason,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some(payload) =
|
||||
maybe_build_stream_plan_payload(state, parts, trace_id, decision, body_json).await?
|
||||
else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
let GatewayControlPlanResponse {
|
||||
@@ -86,17 +90,19 @@ pub(crate) async fn maybe_execute_stream_via_plan_fallback(
|
||||
} = payload;
|
||||
|
||||
let (Some(plan_kind), Some(plan)) = (plan_kind, plan) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
execute_execution_runtime_stream(
|
||||
execute_stream_plan_and_reports(
|
||||
state,
|
||||
plan,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind.as_str(),
|
||||
report_kind,
|
||||
report_context,
|
||||
vec![LocalStreamPlanAndReport {
|
||||
plan,
|
||||
report_kind,
|
||||
report_context,
|
||||
}],
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -4,20 +4,21 @@ use std::collections::BTreeMap;
|
||||
|
||||
use crate::ai_pipeline_api::{
|
||||
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
|
||||
supports_stream_scheduler_decision_kind, OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
supports_stream_scheduler_decision_kind, LocalStreamPlanAndReport,
|
||||
OPENAI_VIDEO_CONTENT_PLAN_KIND,
|
||||
};
|
||||
use crate::api::response::build_client_response_from_parts;
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::execution_runtime::execute_execution_runtime_stream;
|
||||
use crate::{AppState, GatewayError, GatewayFallbackReason};
|
||||
|
||||
use super::{
|
||||
build_direct_plan_bypass_cache_key, maybe_execute_stream_via_local_decision,
|
||||
maybe_execute_stream_via_local_gemini_files_decision,
|
||||
build_direct_plan_bypass_cache_key, execute_stream_plan_and_reports,
|
||||
maybe_execute_stream_via_local_decision, maybe_execute_stream_via_local_gemini_files_decision,
|
||||
maybe_execute_stream_via_local_openai_cli_decision,
|
||||
maybe_execute_stream_via_local_same_format_provider_decision,
|
||||
maybe_execute_stream_via_local_standard_decision, maybe_execute_stream_via_plan_fallback,
|
||||
maybe_execute_stream_via_remote_decision, parse_local_request_body, should_skip_direct_plan,
|
||||
LocalExecutionRequestOutcome,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_execute_via_stream_decision_path(
|
||||
@@ -26,71 +27,96 @@ pub(crate) async fn maybe_execute_via_stream_decision_path(
|
||||
body_bytes: &Bytes,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
let Some(plan_kind) = resolve_execution_runtime_stream_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
let Some((body_json, body_base64)) = parse_local_request_body(parts, body_bytes) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
if !is_matching_stream_request(plan_kind, parts, &body_json) {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
let bypass_cache_key =
|
||||
build_direct_plan_bypass_cache_key(plan_kind, parts, body_bytes, decision);
|
||||
if should_skip_direct_plan(state, &bypass_cache_key) {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
if let Some(response) =
|
||||
maybe_execute_local_video_task_content_stream(state, parts, trace_id, decision, plan_kind)
|
||||
.await?
|
||||
let mut exhausted = None;
|
||||
|
||||
match maybe_execute_local_video_task_content_stream(state, parts, trace_id, decision, plan_kind)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if supports_stream_scheduler_decision_kind(plan_kind) {
|
||||
if let Some(response) = maybe_execute_stream_via_local_decision(
|
||||
match maybe_execute_stream_via_local_decision(
|
||||
state, parts, trace_id, decision, &body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_stream_via_local_openai_cli_decision(
|
||||
match maybe_execute_stream_via_local_openai_cli_decision(
|
||||
state, parts, trace_id, decision, &body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_stream_via_local_standard_decision(
|
||||
match maybe_execute_stream_via_local_standard_decision(
|
||||
state, parts, trace_id, decision, &body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_stream_via_local_same_format_provider_decision(
|
||||
match maybe_execute_stream_via_local_same_format_provider_decision(
|
||||
state, parts, trace_id, decision, &body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_stream_via_local_gemini_files_decision(
|
||||
match maybe_execute_stream_via_local_gemini_files_decision(
|
||||
state, parts, trace_id, decision, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_stream_via_remote_decision(
|
||||
@@ -98,11 +124,11 @@ pub(crate) async fn maybe_execute_via_stream_decision_path(
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
}
|
||||
|
||||
maybe_execute_stream_via_plan_fallback(
|
||||
match maybe_execute_stream_via_plan_fallback(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
@@ -117,7 +143,18 @@ pub(crate) async fn maybe_execute_via_stream_decision_path(
|
||||
GatewayFallbackReason::SchedulerDecisionUnsupported
|
||||
},
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
Ok(LocalExecutionRequestOutcome::Responded(response))
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => {
|
||||
Ok(LocalExecutionRequestOutcome::Exhausted(outcome))
|
||||
}
|
||||
LocalExecutionRequestOutcome::NoPath => Ok(exhausted
|
||||
.map(LocalExecutionRequestOutcome::Exhausted)
|
||||
.unwrap_or(LocalExecutionRequestOutcome::NoPath)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn maybe_execute_local_video_task_content_stream(
|
||||
@@ -126,11 +163,11 @@ async fn maybe_execute_local_video_task_content_stream(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
if plan_kind != OPENAI_VIDEO_CONTENT_PLAN_KIND
|
||||
|| decision.route_family.as_deref() != Some("openai")
|
||||
{
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
let _ = state
|
||||
@@ -155,22 +192,28 @@ async fn maybe_execute_local_video_task_content_stream(
|
||||
parts.uri.query(),
|
||||
trace_id,
|
||||
) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
match action {
|
||||
crate::video_tasks::LocalVideoTaskContentAction::Immediate {
|
||||
status_code,
|
||||
body_json,
|
||||
} => Ok(Some(build_json_response(
|
||||
trace_id,
|
||||
decision,
|
||||
status_code,
|
||||
&body_json,
|
||||
)?)),
|
||||
} => Ok(LocalExecutionRequestOutcome::Responded(
|
||||
build_json_response(trace_id, decision, status_code, &body_json)?,
|
||||
)),
|
||||
crate::video_tasks::LocalVideoTaskContentAction::StreamPlan(plan) => {
|
||||
execute_execution_runtime_stream(
|
||||
state, *plan, trace_id, decision, plan_kind, None, None,
|
||||
let plan = *plan;
|
||||
execute_stream_plan_and_reports(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
vec![LocalStreamPlanAndReport {
|
||||
plan,
|
||||
report_kind: None,
|
||||
report_context: None,
|
||||
}],
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -5,23 +5,22 @@ use std::collections::BTreeMap;
|
||||
use crate::ai_pipeline_api::{
|
||||
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
|
||||
resolve_execution_runtime_sync_plan_kind, supports_sync_scheduler_decision_kind,
|
||||
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
LocalSyncPlanAndReport, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
|
||||
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::api::response::build_client_response_from_parts;
|
||||
use crate::control::resolve_execution_runtime_auth_context;
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::execution_runtime::execute_execution_runtime_sync;
|
||||
use crate::{AppState, GatewayError, GatewayFallbackReason};
|
||||
|
||||
use super::{
|
||||
build_direct_plan_bypass_cache_key, maybe_execute_sync_via_local_decision,
|
||||
maybe_execute_sync_via_local_gemini_files_decision,
|
||||
build_direct_plan_bypass_cache_key, execute_sync_plan_and_reports,
|
||||
maybe_execute_sync_via_local_decision, maybe_execute_sync_via_local_gemini_files_decision,
|
||||
maybe_execute_sync_via_local_openai_cli_decision,
|
||||
maybe_execute_sync_via_local_same_format_provider_decision,
|
||||
maybe_execute_sync_via_local_standard_decision, maybe_execute_sync_via_local_video_decision,
|
||||
maybe_execute_sync_via_plan_fallback, maybe_execute_sync_via_remote_decision,
|
||||
parse_local_request_body, should_skip_direct_plan,
|
||||
parse_local_request_body, should_skip_direct_plan, LocalExecutionRequestOutcome,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_execute_via_sync_decision_path(
|
||||
@@ -30,83 +29,109 @@ pub(crate) async fn maybe_execute_via_sync_decision_path(
|
||||
body_bytes: &Bytes,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
if let Some(response) =
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
if let LocalExecutionRequestOutcome::Responded(response) =
|
||||
maybe_build_local_video_task_read_response(state, parts, trace_id, decision).await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
|
||||
let Some(plan_kind) = resolve_execution_runtime_sync_plan_kind(parts, decision) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
let Some((body_json, body_base64)) = parse_local_request_body(parts, body_bytes) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
if let Some(stream_plan_kind) = resolve_execution_runtime_stream_plan_kind(parts, decision) {
|
||||
if is_matching_stream_request(stream_plan_kind, parts, &body_json) {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
}
|
||||
|
||||
let bypass_cache_key =
|
||||
build_direct_plan_bypass_cache_key(plan_kind, parts, body_bytes, decision);
|
||||
if should_skip_direct_plan(state, &bypass_cache_key) {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_local_video_task_follow_up_sync(
|
||||
let mut exhausted = None;
|
||||
|
||||
match maybe_execute_local_video_task_follow_up_sync(
|
||||
state, parts, &body_json, trace_id, decision, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if supports_sync_scheduler_decision_kind(plan_kind) {
|
||||
if let Some(response) = maybe_execute_sync_via_local_video_decision(
|
||||
match maybe_execute_sync_via_local_video_decision(
|
||||
state, parts, &body_json, trace_id, decision, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_sync_via_local_decision(
|
||||
match maybe_execute_sync_via_local_decision(
|
||||
state, parts, trace_id, decision, &body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_sync_via_local_openai_cli_decision(
|
||||
match maybe_execute_sync_via_local_openai_cli_decision(
|
||||
state, parts, trace_id, decision, &body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_sync_via_local_standard_decision(
|
||||
match maybe_execute_sync_via_local_standard_decision(
|
||||
state, parts, trace_id, decision, &body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_sync_via_local_same_format_provider_decision(
|
||||
match maybe_execute_sync_via_local_same_format_provider_decision(
|
||||
state, parts, trace_id, decision, &body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_sync_via_local_gemini_files_decision(
|
||||
match maybe_execute_sync_via_local_gemini_files_decision(
|
||||
state,
|
||||
parts,
|
||||
&body_json,
|
||||
@@ -118,7 +143,11 @@ pub(crate) async fn maybe_execute_via_sync_decision_path(
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => exhausted = Some(outcome),
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
|
||||
if let Some(response) = maybe_execute_sync_via_remote_decision(
|
||||
@@ -126,11 +155,11 @@ pub(crate) async fn maybe_execute_via_sync_decision_path(
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(response));
|
||||
return Ok(LocalExecutionRequestOutcome::Responded(response));
|
||||
}
|
||||
}
|
||||
|
||||
maybe_execute_sync_via_plan_fallback(
|
||||
match maybe_execute_sync_via_plan_fallback(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
@@ -145,7 +174,18 @@ pub(crate) async fn maybe_execute_via_sync_decision_path(
|
||||
GatewayFallbackReason::SchedulerDecisionUnsupported
|
||||
},
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
LocalExecutionRequestOutcome::Responded(response) => {
|
||||
Ok(LocalExecutionRequestOutcome::Responded(response))
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => {
|
||||
Ok(LocalExecutionRequestOutcome::Exhausted(outcome))
|
||||
}
|
||||
LocalExecutionRequestOutcome::NoPath => Ok(exhausted
|
||||
.map(LocalExecutionRequestOutcome::Exhausted)
|
||||
.unwrap_or(LocalExecutionRequestOutcome::NoPath)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn maybe_build_local_video_task_read_response(
|
||||
@@ -153,9 +193,9 @@ async fn maybe_build_local_video_task_read_response(
|
||||
parts: &http::request::Parts,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
if parts.method != http::Method::GET || decision.route_kind.as_deref() != Some("video") {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
let _ = state
|
||||
@@ -187,7 +227,7 @@ async fn maybe_build_local_video_task_read_response(
|
||||
}
|
||||
};
|
||||
let Some(read_response) = read_response else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
let body_bytes = serde_json::to_vec(&read_response.body_json)
|
||||
@@ -196,13 +236,15 @@ async fn maybe_build_local_video_task_read_response(
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
|
||||
Ok(Some(build_client_response_from_parts(
|
||||
read_response.status_code,
|
||||
&headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?))
|
||||
Ok(LocalExecutionRequestOutcome::Responded(
|
||||
build_client_response_from_parts(
|
||||
read_response.status_code,
|
||||
&headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?,
|
||||
))
|
||||
}
|
||||
|
||||
async fn maybe_execute_local_video_task_follow_up_sync(
|
||||
@@ -212,7 +254,7 @@ async fn maybe_execute_local_video_task_follow_up_sync(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
|
||||
if !matches!(
|
||||
plan_kind,
|
||||
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND
|
||||
@@ -220,7 +262,7 @@ async fn maybe_execute_local_video_task_follow_up_sync(
|
||||
| OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND
|
||||
| GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND
|
||||
) {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
}
|
||||
|
||||
let _ = state
|
||||
@@ -242,18 +284,20 @@ async fn maybe_execute_local_video_task_follow_up_sync(
|
||||
auth_context.as_ref(),
|
||||
trace_id,
|
||||
) else {
|
||||
return Ok(None);
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
execute_execution_runtime_sync(
|
||||
execute_sync_plan_and_reports(
|
||||
state,
|
||||
parts.uri.path(),
|
||||
follow_up.plan,
|
||||
parts,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
follow_up.report_kind,
|
||||
follow_up.report_context,
|
||||
vec![LocalSyncPlanAndReport {
|
||||
plan: follow_up.plan,
|
||||
report_kind: follow_up.report_kind,
|
||||
report_context: follow_up.report_context,
|
||||
}],
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -29,7 +29,10 @@ use crate::control::{
|
||||
should_buffer_request_for_local_auth, trusted_auth_local_rejection, GatewayControlDecision,
|
||||
GatewayPublicRequestContext,
|
||||
};
|
||||
use crate::executor::{maybe_execute_stream_request, maybe_execute_sync_request};
|
||||
use crate::executor::{
|
||||
maybe_execute_stream_request, maybe_execute_sync_request,
|
||||
record_failed_usage_for_exhausted_request, LocalExecutionRequestOutcome,
|
||||
};
|
||||
use crate::handlers::shared::{
|
||||
build_admin_proxy_auth_required_response, build_unhandled_admin_proxy_response,
|
||||
local_proxy_route_requires_buffered_body, request_enables_control_execute,
|
||||
@@ -44,7 +47,6 @@ use crate::{
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::extract::{ConnectInfo, Request, State};
|
||||
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
|
||||
use chrono::Utc;
|
||||
use std::time::Instant;
|
||||
use tracing::{info, warn};
|
||||
|
||||
@@ -645,8 +647,9 @@ pub(crate) async fn proxy_request(
|
||||
.as_ref()
|
||||
.expect("execution runtime/control auth gate should have buffered request body");
|
||||
let stream_request = request_wants_stream(&request_context, buffered_body);
|
||||
let mut local_execution_exhaustion = None;
|
||||
if stream_request {
|
||||
if let Some(execution_runtime_response) = maybe_execute_stream_request(
|
||||
match maybe_execute_stream_request(
|
||||
&state,
|
||||
&parts,
|
||||
buffered_body,
|
||||
@@ -655,35 +658,46 @@ pub(crate) async fn proxy_request(
|
||||
)
|
||||
.await?
|
||||
{
|
||||
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
|
||||
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
|
||||
return Ok(finalize_gateway_response_with_context(
|
||||
&state,
|
||||
execution_runtime_response,
|
||||
&remote_addr,
|
||||
&request_context,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => {
|
||||
local_execution_exhaustion = Some(outcome);
|
||||
}
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
}
|
||||
match maybe_execute_sync_request(&state, &parts, buffered_body, &trace_id, control_decision)
|
||||
.await?
|
||||
{
|
||||
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
|
||||
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
|
||||
return Ok(finalize_gateway_response_with_context(
|
||||
&state,
|
||||
execution_runtime_response,
|
||||
&remote_addr,
|
||||
&request_context,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if let Some(execution_runtime_response) =
|
||||
maybe_execute_sync_request(&state, &parts, buffered_body, &trace_id, control_decision)
|
||||
.await?
|
||||
{
|
||||
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
|
||||
return Ok(finalize_gateway_response_with_context(
|
||||
&state,
|
||||
execution_runtime_response,
|
||||
&remote_addr,
|
||||
&request_context,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => {
|
||||
local_execution_exhaustion = Some(outcome);
|
||||
}
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
if parts.method != http::Method::POST {
|
||||
if let Some(execution_runtime_response) = maybe_execute_stream_request(
|
||||
match maybe_execute_stream_request(
|
||||
&state,
|
||||
&parts,
|
||||
buffered_body,
|
||||
@@ -692,20 +706,26 @@ pub(crate) async fn proxy_request(
|
||||
)
|
||||
.await?
|
||||
{
|
||||
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
|
||||
return Ok(finalize_gateway_response_with_context(
|
||||
&state,
|
||||
execution_runtime_response,
|
||||
&remote_addr,
|
||||
&request_context,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
|
||||
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
|
||||
return Ok(finalize_gateway_response_with_context(
|
||||
&state,
|
||||
execution_runtime_response,
|
||||
&remote_addr,
|
||||
&request_context,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => {
|
||||
local_execution_exhaustion = Some(outcome);
|
||||
}
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
}
|
||||
if allow_control_execute_fallback {
|
||||
if let Some(control_response) = maybe_execute_via_control(
|
||||
match maybe_execute_via_control(
|
||||
&state,
|
||||
&parts,
|
||||
buffered_body.clone(),
|
||||
@@ -715,41 +735,47 @@ pub(crate) async fn proxy_request(
|
||||
)
|
||||
.await?
|
||||
{
|
||||
let reason = GatewayFallbackReason::ControlExecuteEmergency;
|
||||
let control_execution_path = if stream_request {
|
||||
EXECUTION_PATH_CONTROL_EXECUTE_STREAM
|
||||
} else {
|
||||
EXECUTION_PATH_CONTROL_EXECUTE_SYNC
|
||||
};
|
||||
state.record_fallback_metric(
|
||||
GatewayFallbackMetricKind::ControlExecuteFallback,
|
||||
control_decision,
|
||||
None,
|
||||
Some(control_execution_path),
|
||||
reason,
|
||||
);
|
||||
state.record_fallback_metric(
|
||||
GatewayFallbackMetricKind::RemoteExecuteEmergency,
|
||||
control_decision,
|
||||
None,
|
||||
Some(control_execution_path),
|
||||
reason,
|
||||
);
|
||||
let mut control_response = control_response;
|
||||
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
|
||||
control_response.headers_mut().insert(
|
||||
HeaderName::from_static(DEPENDENCY_REASON_HEADER),
|
||||
HeaderValue::from_static(reason.as_label_value()),
|
||||
);
|
||||
return Ok(finalize_gateway_response_with_context(
|
||||
&state,
|
||||
control_response,
|
||||
&remote_addr,
|
||||
&request_context,
|
||||
control_execution_path,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
LocalExecutionRequestOutcome::Responded(control_response) => {
|
||||
let reason = GatewayFallbackReason::ControlExecuteEmergency;
|
||||
let control_execution_path = if stream_request {
|
||||
EXECUTION_PATH_CONTROL_EXECUTE_STREAM
|
||||
} else {
|
||||
EXECUTION_PATH_CONTROL_EXECUTE_SYNC
|
||||
};
|
||||
state.record_fallback_metric(
|
||||
GatewayFallbackMetricKind::ControlExecuteFallback,
|
||||
control_decision,
|
||||
None,
|
||||
Some(control_execution_path),
|
||||
reason,
|
||||
);
|
||||
state.record_fallback_metric(
|
||||
GatewayFallbackMetricKind::RemoteExecuteEmergency,
|
||||
control_decision,
|
||||
None,
|
||||
Some(control_execution_path),
|
||||
reason,
|
||||
);
|
||||
let mut control_response = control_response;
|
||||
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
|
||||
control_response.headers_mut().insert(
|
||||
HeaderName::from_static(DEPENDENCY_REASON_HEADER),
|
||||
HeaderValue::from_static(reason.as_label_value()),
|
||||
);
|
||||
return Ok(finalize_gateway_response_with_context(
|
||||
&state,
|
||||
control_response,
|
||||
&remote_addr,
|
||||
&request_context,
|
||||
control_execution_path,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
}
|
||||
LocalExecutionRequestOutcome::Exhausted(outcome) => {
|
||||
local_execution_exhaustion = Some(outcome);
|
||||
}
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
}
|
||||
let local_execution_runtime_miss_detail =
|
||||
@@ -779,6 +805,16 @@ pub(crate) async fn proxy_request(
|
||||
"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_diagnostic.as_ref(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
let mut response = build_local_http_error_response(
|
||||
&trace_id,
|
||||
control_decision,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use aether_billing::normalize_input_tokens_for_billing;
|
||||
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageAuditListQuery};
|
||||
use axum::{
|
||||
body::Body,
|
||||
@@ -88,11 +89,33 @@ fn parse_users_me_usage_ids(query: Option<&str>) -> Option<BTreeSet<String>> {
|
||||
(!values.is_empty()).then_some(values)
|
||||
}
|
||||
|
||||
fn users_me_usage_cache_creation_tokens(item: &StoredRequestUsageAudit) -> u64 {
|
||||
let classified = item
|
||||
.cache_creation_ephemeral_5m_input_tokens
|
||||
.saturating_add(item.cache_creation_ephemeral_1h_input_tokens);
|
||||
if item.cache_creation_input_tokens == 0 && classified > 0 {
|
||||
classified
|
||||
} else {
|
||||
item.cache_creation_input_tokens
|
||||
}
|
||||
}
|
||||
|
||||
fn users_me_usage_total_input_context(item: &StoredRequestUsageAudit) -> u64 {
|
||||
item.input_tokens
|
||||
.saturating_add(users_me_usage_cache_creation_tokens(item))
|
||||
.saturating_add(item.cache_read_input_tokens)
|
||||
}
|
||||
|
||||
fn users_me_usage_effective_input_tokens(item: &StoredRequestUsageAudit) -> u64 {
|
||||
let api_format = item
|
||||
.endpoint_api_format
|
||||
.as_deref()
|
||||
.or(item.api_format.as_deref());
|
||||
let input_tokens = i64::try_from(item.input_tokens).unwrap_or(i64::MAX);
|
||||
let cache_read_tokens = i64::try_from(item.cache_read_input_tokens).unwrap_or(i64::MAX);
|
||||
normalize_input_tokens_for_billing(api_format, input_tokens, cache_read_tokens) as u64
|
||||
}
|
||||
|
||||
fn users_me_usage_effective_unix_secs(item: &StoredRequestUsageAudit) -> u64 {
|
||||
item.finalized_at_unix_secs
|
||||
.unwrap_or(item.created_at_unix_ms)
|
||||
@@ -152,6 +175,7 @@ fn build_users_me_usage_record_payload(
|
||||
"endpoint_api_format": item.endpoint_api_format,
|
||||
"has_format_conversion": item.has_format_conversion,
|
||||
"input_tokens": item.input_tokens,
|
||||
"effective_input_tokens": users_me_usage_effective_input_tokens(item),
|
||||
"output_tokens": item.output_tokens,
|
||||
"total_tokens": item.total_tokens,
|
||||
"cost": round_to(item.total_cost_usd, 6),
|
||||
@@ -161,6 +185,8 @@ fn build_users_me_usage_record_payload(
|
||||
"status": item.status,
|
||||
"created_at": unix_secs_to_rfc3339(item.created_at_unix_ms),
|
||||
"cache_creation_input_tokens": item.cache_creation_input_tokens,
|
||||
"cache_creation_ephemeral_5m_input_tokens": item.cache_creation_ephemeral_5m_input_tokens,
|
||||
"cache_creation_ephemeral_1h_input_tokens": item.cache_creation_ephemeral_1h_input_tokens,
|
||||
"cache_read_input_tokens": item.cache_read_input_tokens,
|
||||
"status_code": item.status_code,
|
||||
"error_message": item.error_message,
|
||||
@@ -186,8 +212,11 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
|
||||
"id": item.id,
|
||||
"status": item.status,
|
||||
"input_tokens": item.input_tokens,
|
||||
"effective_input_tokens": users_me_usage_effective_input_tokens(item),
|
||||
"output_tokens": item.output_tokens,
|
||||
"cache_creation_input_tokens": item.cache_creation_input_tokens,
|
||||
"cache_creation_ephemeral_5m_input_tokens": item.cache_creation_ephemeral_5m_input_tokens,
|
||||
"cache_creation_ephemeral_1h_input_tokens": item.cache_creation_ephemeral_1h_input_tokens,
|
||||
"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),
|
||||
@@ -231,10 +260,13 @@ fn build_users_me_usage_summary_by_model(
|
||||
"model": item.model,
|
||||
"requests": 0_u64,
|
||||
"input_tokens": 0_u64,
|
||||
"effective_input_tokens": 0_u64,
|
||||
"output_tokens": 0_u64,
|
||||
"total_tokens": 0_u64,
|
||||
"cache_read_tokens": 0_u64,
|
||||
"cache_creation_tokens": 0_u64,
|
||||
"cache_creation_ephemeral_5m_tokens": 0_u64,
|
||||
"cache_creation_ephemeral_1h_tokens": 0_u64,
|
||||
"total_input_context": 0_u64,
|
||||
"cache_hit_rate": 0.0,
|
||||
"total_cost_usd": 0.0,
|
||||
@@ -245,6 +277,10 @@ fn build_users_me_usage_summary_by_model(
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(item.input_tokens));
|
||||
entry["effective_input_tokens"] = json!(entry["effective_input_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(users_me_usage_effective_input_tokens(item)));
|
||||
entry["output_tokens"] = json!(entry["output_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
@@ -260,7 +296,17 @@ fn build_users_me_usage_summary_by_model(
|
||||
entry["cache_creation_tokens"] = json!(entry["cache_creation_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(item.cache_creation_input_tokens));
|
||||
.saturating_add(users_me_usage_cache_creation_tokens(item)));
|
||||
entry["cache_creation_ephemeral_5m_tokens"] = json!(entry
|
||||
["cache_creation_ephemeral_5m_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(item.cache_creation_ephemeral_5m_input_tokens));
|
||||
entry["cache_creation_ephemeral_1h_tokens"] = json!(entry
|
||||
["cache_creation_ephemeral_1h_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(item.cache_creation_ephemeral_1h_input_tokens));
|
||||
entry["total_input_context"] = json!(entry["total_input_context"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
@@ -320,11 +366,14 @@ fn build_users_me_usage_summary_by_provider(
|
||||
json!({
|
||||
"provider": item.provider_name,
|
||||
"requests": 0_u64,
|
||||
"effective_input_tokens": 0_u64,
|
||||
"total_tokens": 0_u64,
|
||||
"total_input_context": 0_u64,
|
||||
"output_tokens": 0_u64,
|
||||
"cache_read_tokens": 0_u64,
|
||||
"cache_creation_tokens": 0_u64,
|
||||
"cache_creation_ephemeral_5m_tokens": 0_u64,
|
||||
"cache_creation_ephemeral_1h_tokens": 0_u64,
|
||||
"cache_hit_rate": 0.0,
|
||||
"total_cost_usd": 0.0,
|
||||
"success_rate": 0.0,
|
||||
@@ -343,6 +392,10 @@ fn build_users_me_usage_summary_by_provider(
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(users_me_usage_total_input_context(item)));
|
||||
entry["effective_input_tokens"] = json!(entry["effective_input_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(users_me_usage_effective_input_tokens(item)));
|
||||
entry["output_tokens"] = json!(entry["output_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
@@ -354,7 +407,17 @@ fn build_users_me_usage_summary_by_provider(
|
||||
entry["cache_creation_tokens"] = json!(entry["cache_creation_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(item.cache_creation_input_tokens));
|
||||
.saturating_add(users_me_usage_cache_creation_tokens(item)));
|
||||
entry["cache_creation_ephemeral_5m_tokens"] = json!(entry
|
||||
["cache_creation_ephemeral_5m_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(item.cache_creation_ephemeral_5m_input_tokens));
|
||||
entry["cache_creation_ephemeral_1h_tokens"] = json!(entry
|
||||
["cache_creation_ephemeral_1h_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(item.cache_creation_ephemeral_1h_input_tokens));
|
||||
entry["total_cost_usd"] =
|
||||
json!(entry["total_cost_usd"].as_f64().unwrap_or(0.0) + item.total_cost_usd);
|
||||
|
||||
@@ -445,10 +508,13 @@ fn build_users_me_usage_summary_by_api_format(
|
||||
"api_format": api_format,
|
||||
"request_count": 0_u64,
|
||||
"total_tokens": 0_u64,
|
||||
"effective_input_tokens": 0_u64,
|
||||
"total_input_context": 0_u64,
|
||||
"output_tokens": 0_u64,
|
||||
"cache_read_tokens": 0_u64,
|
||||
"cache_creation_tokens": 0_u64,
|
||||
"cache_creation_ephemeral_5m_tokens": 0_u64,
|
||||
"cache_creation_ephemeral_1h_tokens": 0_u64,
|
||||
"cache_hit_rate": 0.0,
|
||||
"total_cost_usd": 0.0,
|
||||
"avg_response_time_ms": 0.0,
|
||||
@@ -468,6 +534,10 @@ fn build_users_me_usage_summary_by_api_format(
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(users_me_usage_total_input_context(item)));
|
||||
entry["effective_input_tokens"] = json!(entry["effective_input_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(users_me_usage_effective_input_tokens(item)));
|
||||
entry["output_tokens"] = json!(entry["output_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
@@ -479,7 +549,17 @@ fn build_users_me_usage_summary_by_api_format(
|
||||
entry["cache_creation_tokens"] = json!(entry["cache_creation_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(item.cache_creation_input_tokens));
|
||||
.saturating_add(users_me_usage_cache_creation_tokens(item)));
|
||||
entry["cache_creation_ephemeral_5m_tokens"] = json!(entry
|
||||
["cache_creation_ephemeral_5m_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(item.cache_creation_ephemeral_5m_input_tokens));
|
||||
entry["cache_creation_ephemeral_1h_tokens"] = json!(entry
|
||||
["cache_creation_ephemeral_1h_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(item.cache_creation_ephemeral_1h_input_tokens));
|
||||
entry["total_cost_usd"] =
|
||||
json!(entry["total_cost_usd"].as_f64().unwrap_or(0.0) + item.total_cost_usd);
|
||||
if let Some(response_time_ms) = item.response_time_ms {
|
||||
|
||||
@@ -18,6 +18,8 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy
|
||||
url: String,
|
||||
model: String,
|
||||
auth_header_value: String,
|
||||
anthropic_version: String,
|
||||
anthropic_beta: String,
|
||||
endpoint_tag: String,
|
||||
metadata_mode: String,
|
||||
metadata_source: String,
|
||||
@@ -261,6 +263,18 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
anthropic_version: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("anthropic-version"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
anthropic_beta: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("anthropic-beta"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
@@ -357,6 +371,8 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy
|
||||
.post(format!("{gateway_url}/v1/messages"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("x-api-key", "sk-client-claude-chat-local")
|
||||
.header("anthropic-version", "2023-06-01")
|
||||
.header("anthropic-beta", "prompt-caching-2024-07-31,context-1m-2025-08-07")
|
||||
.header(TRACE_ID_HEADER, "trace-claude-chat-local-123")
|
||||
.body(
|
||||
"{\"model\":\"claude-sonnet-4-5\",\"messages\":[],\"metadata\":{\"client\":\"desktop-claude\"}}",
|
||||
@@ -390,6 +406,14 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy
|
||||
seen_execution_runtime_request.auth_header_value,
|
||||
"sk-upstream-claude-chat"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.anthropic_version,
|
||||
"2023-06-01"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.anthropic_beta,
|
||||
"prompt-caching-2024-07-31,context-1m-2025-08-07"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"claude-chat-local"
|
||||
|
||||
@@ -18,6 +18,8 @@ async fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_syn
|
||||
url: String,
|
||||
model: String,
|
||||
authorization: String,
|
||||
anthropic_version: String,
|
||||
anthropic_beta: String,
|
||||
endpoint_tag: String,
|
||||
metadata_mode: String,
|
||||
metadata_source: String,
|
||||
@@ -261,6 +263,18 @@ async fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_syn
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
anthropic_version: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("anthropic-version"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
anthropic_beta: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("anthropic-beta"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
@@ -359,6 +373,8 @@ async fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_syn
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-claude-cli-local",
|
||||
)
|
||||
.header("anthropic-version", "2023-06-01")
|
||||
.header("anthropic-beta", "prompt-caching-2024-07-31")
|
||||
.header(TRACE_ID_HEADER, "trace-claude-cli-local-sync-123")
|
||||
.body(
|
||||
"{\"model\":\"claude-code\",\"messages\":[],\"metadata\":{\"client\":\"desktop-claude-cli\"}}",
|
||||
@@ -398,6 +414,14 @@ async fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_syn
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer sk-upstream-claude-cli"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.anthropic_version,
|
||||
"2023-06-01"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.anthropic_beta,
|
||||
"prompt-caching-2024-07-31"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"claude-cli-local"
|
||||
|
||||
@@ -151,7 +151,7 @@ fn sample_usage_row(
|
||||
actual_total_cost_usd: f64,
|
||||
created_at_unix_ms: i64,
|
||||
) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
let mut usage = StoredRequestUsageAudit::new(
|
||||
id.to_string(),
|
||||
request_id.to_string(),
|
||||
user_id.map(str::to_string),
|
||||
@@ -190,7 +190,10 @@ fn sample_usage_row(
|
||||
Some(created_at_unix_ms + 2),
|
||||
)
|
||||
.expect("usage row should build")
|
||||
.with_cache_input_tokens(15, 5)
|
||||
.with_cache_input_tokens(15, 5);
|
||||
usage.cache_creation_ephemeral_5m_input_tokens = 6;
|
||||
usage.cache_creation_ephemeral_1h_input_tokens = 9;
|
||||
usage
|
||||
}
|
||||
|
||||
fn sample_user_summary(id: &str, username: &str) -> StoredUserSummary {
|
||||
@@ -272,6 +275,14 @@ async fn gateway_handles_admin_usage_stats_locally_with_trusted_admin_principal(
|
||||
assert_eq!(payload["total_tokens"], 240);
|
||||
assert_eq!(payload["error_count"], 1);
|
||||
assert_eq!(payload["cache_stats"]["cache_creation_tokens"], 30);
|
||||
assert_eq!(
|
||||
payload["cache_stats"]["cache_creation_ephemeral_5m_tokens"],
|
||||
12
|
||||
);
|
||||
assert_eq!(
|
||||
payload["cache_stats"]["cache_creation_ephemeral_1h_tokens"],
|
||||
18
|
||||
);
|
||||
assert_eq!(payload["cache_stats"]["cache_read_tokens"], 10);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -296,52 +307,66 @@ async fn gateway_handles_admin_usage_aggregation_stats_locally_with_trusted_admi
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_usage_upstream("/api/admin/usage/aggregation/stats").await;
|
||||
|
||||
let mut usage_1 = sample_usage_row(
|
||||
"usage-1",
|
||||
"req-1",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
120,
|
||||
30,
|
||||
0.3,
|
||||
0.36,
|
||||
DAY_1_UNIX_SECS,
|
||||
);
|
||||
usage_1.provider_id = Some("provider-openai".to_string());
|
||||
usage_1.total_tokens = usage_1.input_tokens;
|
||||
|
||||
let mut usage_2 = sample_usage_row(
|
||||
"usage-2",
|
||||
"req-2",
|
||||
Some("user-2"),
|
||||
Some("key-2"),
|
||||
Some("secondary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
40,
|
||||
10,
|
||||
0.1,
|
||||
0.12,
|
||||
DAY_2_UNIX_SECS,
|
||||
);
|
||||
usage_2.provider_id = Some("provider-openai".to_string());
|
||||
usage_2.total_tokens = usage_2.input_tokens;
|
||||
|
||||
let mut usage_3 = sample_usage_row(
|
||||
"usage-3",
|
||||
"req-3",
|
||||
Some("user-2"),
|
||||
Some("key-2"),
|
||||
Some("secondary"),
|
||||
"Anthropic",
|
||||
"claude-3-7",
|
||||
"completed",
|
||||
60,
|
||||
20,
|
||||
0.2,
|
||||
0.24,
|
||||
DAY_2_UNIX_SECS,
|
||||
);
|
||||
usage_3.provider_id = Some("provider-anthropic".to_string());
|
||||
usage_3.total_tokens = usage_3.input_tokens;
|
||||
usage_3.api_format = Some("claude:cli".to_string());
|
||||
usage_3.api_family = Some("claude".to_string());
|
||||
usage_3.endpoint_api_format = Some("claude:cli".to_string());
|
||||
usage_3.provider_api_family = Some("claude".to_string());
|
||||
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_usage_row(
|
||||
"usage-1",
|
||||
"req-1",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
120,
|
||||
30,
|
||||
0.3,
|
||||
0.36,
|
||||
DAY_1_UNIX_SECS,
|
||||
),
|
||||
sample_usage_row(
|
||||
"usage-2",
|
||||
"req-2",
|
||||
Some("user-2"),
|
||||
Some("key-2"),
|
||||
Some("secondary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
40,
|
||||
10,
|
||||
0.1,
|
||||
0.12,
|
||||
DAY_2_UNIX_SECS,
|
||||
),
|
||||
sample_usage_row(
|
||||
"usage-3",
|
||||
"req-3",
|
||||
Some("user-2"),
|
||||
Some("key-2"),
|
||||
Some("secondary"),
|
||||
"Anthropic",
|
||||
"claude-3-7",
|
||||
"completed",
|
||||
60,
|
||||
20,
|
||||
0.2,
|
||||
0.24,
|
||||
DAY_2_UNIX_SECS,
|
||||
),
|
||||
usage_1, usage_2, usage_3,
|
||||
]));
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
@@ -366,7 +391,53 @@ async fn gateway_handles_admin_usage_aggregation_stats_locally_with_trusted_admi
|
||||
assert_eq!(items.len(), 2);
|
||||
assert_eq!(items[0]["model"], "gpt-5");
|
||||
assert_eq!(items[0]["request_count"], 2);
|
||||
assert_eq!(items[0]["output_tokens"], 40);
|
||||
assert_eq!(items[0]["effective_input_tokens"], 150);
|
||||
assert_eq!(items[0]["total_input_context"], 200);
|
||||
assert_eq!(items[0]["cache_creation_tokens"], 30);
|
||||
assert_eq!(items[0]["cache_creation_ephemeral_5m_tokens"], 12);
|
||||
assert_eq!(items[0]["cache_creation_ephemeral_1h_tokens"], 18);
|
||||
assert_eq!(items[0]["cache_hit_rate"], 5.0);
|
||||
assert_eq!(items[1]["model"], "claude-3-7");
|
||||
assert_eq!(items[1]["output_tokens"], 20);
|
||||
|
||||
let provider_response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/usage/aggregation/stats?group_by=provider&limit=10"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(provider_response.status(), StatusCode::OK);
|
||||
let provider_payload: serde_json::Value = provider_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
let provider_items = provider_payload.as_array().expect("array response");
|
||||
assert_eq!(provider_items.len(), 2);
|
||||
assert_eq!(provider_items[0]["provider"], "OpenAI");
|
||||
assert_eq!(provider_items[0]["output_tokens"], 40);
|
||||
assert_eq!(provider_items[1]["provider"], "Anthropic");
|
||||
assert_eq!(provider_items[1]["output_tokens"], 20);
|
||||
|
||||
let api_format_response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/usage/aggregation/stats?group_by=api_format&limit=10"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(api_format_response.status(), StatusCode::OK);
|
||||
let api_format_payload: serde_json::Value = api_format_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
let api_format_items = api_format_payload.as_array().expect("array response");
|
||||
assert_eq!(api_format_items.len(), 2);
|
||||
assert_eq!(api_format_items[0]["api_format"], "openai:chat");
|
||||
assert_eq!(api_format_items[0]["output_tokens"], 40);
|
||||
assert_eq!(api_format_items[1]["api_format"], "claude:cli");
|
||||
assert_eq!(api_format_items[1]["output_tokens"], 20);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
@@ -616,6 +687,7 @@ async fn gateway_handles_admin_usage_active_locally_with_trusted_admin_principal
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["requests"].as_array().expect("array").len(), 1);
|
||||
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]["provider_key_name"],
|
||||
@@ -707,6 +779,7 @@ async fn gateway_handles_admin_usage_records_locally_with_trusted_admin_principa
|
||||
payload["records"][0]["provider_key_name"],
|
||||
"upstream-primary"
|
||||
);
|
||||
assert_eq!(payload["records"][0]["effective_input_tokens"], 35);
|
||||
assert_eq!(payload["records"][0]["first_byte_time_ms"], 120);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -870,6 +943,7 @@ async fn gateway_handles_admin_usage_detail_locally_with_trusted_admin_principal
|
||||
assert_eq!(payload["api_key"]["name"], "primary");
|
||||
assert_eq!(payload["provider"], "OpenAI");
|
||||
assert_eq!(payload["model"], "gpt-5");
|
||||
assert_eq!(payload["effective_input_tokens"], 115);
|
||||
assert_eq!(payload["total_tokens"], 170);
|
||||
assert_eq!(payload["cache_creation_cost"], 0.0);
|
||||
assert_eq!(payload["cache_read_cost"], 0.0);
|
||||
|
||||
@@ -2387,6 +2387,8 @@ fn sample_user_usage_audit(
|
||||
)
|
||||
.expect("usage should build");
|
||||
usage.cache_creation_input_tokens = 10;
|
||||
usage.cache_creation_ephemeral_5m_input_tokens = 4;
|
||||
usage.cache_creation_ephemeral_1h_input_tokens = 6;
|
||||
usage.cache_read_input_tokens = 15;
|
||||
usage.cache_creation_cost_usd = 0.05;
|
||||
usage.cache_read_cost_usd = 0.02;
|
||||
@@ -4733,6 +4735,15 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() {
|
||||
payload["records"].as_array().expect("records array").len(),
|
||||
3
|
||||
);
|
||||
assert_eq!(
|
||||
payload["records"][0]["cache_creation_ephemeral_5m_input_tokens"],
|
||||
4
|
||||
);
|
||||
assert_eq!(payload["records"][0]["effective_input_tokens"], 105);
|
||||
assert_eq!(
|
||||
payload["records"][0]["cache_creation_ephemeral_1h_input_tokens"],
|
||||
6
|
||||
);
|
||||
assert_eq!(
|
||||
payload["summary_by_model"]
|
||||
.as_array()
|
||||
@@ -4740,6 +4751,19 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() {
|
||||
.len(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
payload["summary_by_model"][0]["cache_creation_ephemeral_5m_tokens"],
|
||||
4
|
||||
);
|
||||
assert_eq!(
|
||||
payload["summary_by_model"][0]["cache_creation_ephemeral_1h_tokens"],
|
||||
6
|
||||
);
|
||||
assert_eq!(
|
||||
payload["summary_by_model"][0]["effective_input_tokens"],
|
||||
105
|
||||
);
|
||||
assert_eq!(payload["summary_by_model"][0]["total_input_context"], 145);
|
||||
assert_eq!(payload["billing"]["id"], "wallet-auth-1");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -4822,6 +4846,8 @@ 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]["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);
|
||||
|
||||
gateway_handle.abort();
|
||||
|
||||
@@ -183,6 +183,212 @@ async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exhaust_after_retryable_sync_failure(
|
||||
) {
|
||||
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/chat/completions",
|
||||
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().route(
|
||||
"/v1/execute/sync",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"request_id": "trace-openai-chat-local-report-sync-failure-123",
|
||||
"status_code": 503,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"error": {
|
||||
"message": "primary unavailable"
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 25
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-local-report-sync-failure")),
|
||||
sample_local_openai_auth_snapshot(
|
||||
"api-key-openai-usage-local-failure-1",
|
||||
"user-openai-usage-local-failure-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 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-failure",
|
||||
)
|
||||
.header(
|
||||
TRACE_ID_HEADER,
|
||||
"trace-openai-chat-local-report-sync-failure-123",
|
||||
)
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should complete");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let body_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(body_json["error"]["type"], "http_error");
|
||||
|
||||
let mut stored_usage = None;
|
||||
for _ in 0..50 {
|
||||
stored_usage = usage_repository
|
||||
.find_by_request_id("trace-openai-chat-local-report-sync-failure-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-openai-usage-local-failure-1")
|
||||
);
|
||||
assert_eq!(stored_usage.provider_name, "openai");
|
||||
assert_eq!(stored_usage.model, "gpt-5");
|
||||
assert_eq!(stored_usage.api_format.as_deref(), Some("openai:chat"));
|
||||
assert_eq!(
|
||||
stored_usage
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("trace_id"))
|
||||
.and_then(|value| value.as_str()),
|
||||
Some("trace-openai-chat-local-report-sync-failure-123")
|
||||
);
|
||||
assert_eq!(
|
||||
stored_usage
|
||||
.response_body
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("error"))
|
||||
.and_then(|value| value.get("type"))
|
||||
.and_then(|value| value.as_str()),
|
||||
Some("upstream_error")
|
||||
);
|
||||
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-openai-chat-local-report-sync-failure-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
|
||||
assert_eq!(stored_candidates[0].status_code, Some(503));
|
||||
|
||||
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(
|
||||
) {
|
||||
|
||||
@@ -20,7 +20,7 @@ base64 = "0.22"
|
||||
clap = { version = "4", features = ["derive", "env"] }
|
||||
tracing = "0.1"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
serde_json.workspace = true
|
||||
thiserror = "2"
|
||||
bytes = "1"
|
||||
sha2 = "0.10"
|
||||
|
||||
Reference in New Issue
Block a user