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:
fawney19
2026-04-10 17:44:55 +08:00
parent 5014e2f5fd
commit 010ab127e2
64 changed files with 4217 additions and 477 deletions

View File

@@ -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(),

View File

@@ -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(),

View File

@@ -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!({

View File

@@ -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(),

View File

@@ -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,

View File

@@ -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(),

View File

@@ -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")
);
}
}

View File

@@ -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")
);
}
}

View File

@@ -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> {

View File

@@ -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>)

View File

@@ -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,
};

View File

@@ -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
}

View 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()),
)]))
}

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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,

View File

@@ -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 {

View File

@@ -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"

View File

@@ -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"

View File

@@ -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);

View File

@@ -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();

View File

@@ -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(
) {

View File

@@ -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"