mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 02:47:45 +08:00
refactor(gateway): 重构 ai pipeline 规划链路
This commit is contained in:
fawney19
@@ -8,9 +8,10 @@ use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
strip_deprecated_usage_display_fields, StoredProviderApiKeyUsageSummary,
|
||||
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
|
||||
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageWriteRepository,
|
||||
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
|
||||
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
|
||||
StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository,
|
||||
UsageWriteRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -450,6 +451,12 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
|
||||
if existing.is_some_and(|existing| {
|
||||
usage_status_is_finalized(existing.status.as_str())
|
||||
&& usage_status_is_lifecycle(usage.status.as_str())
|
||||
&& !usage_can_recover_terminal_failure(
|
||||
existing.status.as_str(),
|
||||
existing.billing_status.as_str(),
|
||||
usage.status.as_str(),
|
||||
usage.billing_status.as_str(),
|
||||
)
|
||||
}) {
|
||||
return Ok(existing.expect("existing usage should be present").clone());
|
||||
}
|
||||
@@ -856,6 +863,166 @@ mod tests {
|
||||
assert_eq!(stored.finalized_at_unix_secs, Some(101));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_allows_streaming_recovery_after_void_failure() {
|
||||
let repository = InMemoryUsageReadRepository::default();
|
||||
repository
|
||||
.upsert(UpsertUsageRecord {
|
||||
request_id: "req-recover-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("api-key-1".to_string()),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
target_model: None,
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||||
provider_api_key_id: Some("provider-key-1".to_string()),
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
provider_api_family: Some("openai".to_string()),
|
||||
provider_endpoint_kind: Some("chat".to_string()),
|
||||
has_format_conversion: Some(false),
|
||||
is_stream: Some(false),
|
||||
input_tokens: None,
|
||||
output_tokens: None,
|
||||
total_tokens: None,
|
||||
cache_creation_input_tokens: None,
|
||||
cache_creation_ephemeral_5m_input_tokens: None,
|
||||
cache_creation_ephemeral_1h_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_cost_usd: None,
|
||||
cache_read_cost_usd: None,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd: Some(0.0),
|
||||
actual_total_cost_usd: Some(0.0),
|
||||
status_code: Some(503),
|
||||
error_message: Some("provider timeout".to_string()),
|
||||
error_category: Some("provider_error".to_string()),
|
||||
response_time_ms: Some(90),
|
||||
first_byte_time_ms: None,
|
||||
status: "failed".to_string(),
|
||||
billing_status: "void".to_string(),
|
||||
request_headers: None,
|
||||
request_body: None,
|
||||
request_body_ref: None,
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
provider_request_body_ref: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
response_body_ref: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
client_response_body_ref: None,
|
||||
candidate_id: None,
|
||||
candidate_index: None,
|
||||
key_name: None,
|
||||
planner_kind: None,
|
||||
route_family: None,
|
||||
route_kind: None,
|
||||
execution_path: None,
|
||||
local_execution_runtime_miss_reason: None,
|
||||
request_metadata: None,
|
||||
finalized_at_unix_secs: Some(101),
|
||||
created_at_unix_ms: Some(100),
|
||||
updated_at_unix_secs: 101,
|
||||
})
|
||||
.await
|
||||
.expect("failed usage should upsert");
|
||||
|
||||
repository
|
||||
.upsert(UpsertUsageRecord {
|
||||
request_id: "req-recover-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("api-key-1".to_string()),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
target_model: Some("gpt-5-mini".to_string()),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||||
provider_api_key_id: Some("provider-key-1".to_string()),
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
provider_api_family: Some("openai".to_string()),
|
||||
provider_endpoint_kind: Some("chat".to_string()),
|
||||
has_format_conversion: Some(true),
|
||||
is_stream: Some(true),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: None,
|
||||
total_tokens: None,
|
||||
cache_creation_input_tokens: None,
|
||||
cache_creation_ephemeral_5m_input_tokens: None,
|
||||
cache_creation_ephemeral_1h_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_cost_usd: None,
|
||||
cache_read_cost_usd: None,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd: None,
|
||||
actual_total_cost_usd: None,
|
||||
status_code: None,
|
||||
error_message: None,
|
||||
error_category: None,
|
||||
response_time_ms: Some(45),
|
||||
first_byte_time_ms: Some(12),
|
||||
status: "streaming".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
request_headers: None,
|
||||
request_body: None,
|
||||
request_body_ref: None,
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
provider_request_body_ref: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
response_body_ref: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
client_response_body_ref: None,
|
||||
candidate_id: Some("cand-1".to_string()),
|
||||
candidate_index: Some(1),
|
||||
key_name: Some("primary".to_string()),
|
||||
planner_kind: Some("claude_cli_sync".to_string()),
|
||||
route_family: Some("claude".to_string()),
|
||||
route_kind: Some("cli".to_string()),
|
||||
execution_path: Some("remote".to_string()),
|
||||
local_execution_runtime_miss_reason: None,
|
||||
request_metadata: Some(json!({
|
||||
"trace_id": "trace-recovered"
|
||||
})),
|
||||
finalized_at_unix_secs: None,
|
||||
created_at_unix_ms: Some(100),
|
||||
updated_at_unix_secs: 102,
|
||||
})
|
||||
.await
|
||||
.expect("recovery usage should upsert");
|
||||
|
||||
let stored = repository
|
||||
.find_by_request_id("req-recover-1")
|
||||
.await
|
||||
.expect("usage lookup should succeed")
|
||||
.expect("usage should exist");
|
||||
assert_eq!(stored.status, "streaming");
|
||||
assert_eq!(stored.billing_status, "pending");
|
||||
assert_eq!(stored.status_code, None);
|
||||
assert_eq!(stored.error_message, None);
|
||||
assert_eq!(stored.finalized_at_unix_secs, None);
|
||||
assert_eq!(
|
||||
stored.request_metadata,
|
||||
Some(json!({ "trace_id": "trace-recovered" }))
|
||||
);
|
||||
assert_eq!(stored.total_tokens, 10);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn seed_hydrates_legacy_body_ref_metadata_into_typed_fields() {
|
||||
let repository = InMemoryUsageReadRepository::seed(vec![StoredRequestUsageAudit {
|
||||
|
||||
Reference in New Issue
Block a user