refactor: 抽离 AI pipeline 与调度共享能力逻辑

This commit is contained in:
fawney19
2026-04-10 01:46:14 +08:00
parent b901a6ffc7
commit 5014e2f5fd
255 changed files with 15057 additions and 3115 deletions

View File

@@ -282,7 +282,7 @@ async fn gateway_handles_admin_billing_rule_routes_locally_with_trusted_admin_pr
"input_tokens": { "source": "dimension", "dimension": "input_tokens" }
}),
is_enabled: true,
created_at_unix_secs: 1_710_000_000,
created_at_unix_ms: 1_710_000_000,
updated_at_unix_secs: 1_710_000_100,
},
crate::AdminBillingRuleRecord {
@@ -297,7 +297,7 @@ async fn gateway_handles_admin_billing_rule_routes_locally_with_trusted_admin_pr
"images": { "source": "dimension", "dimension": "image_count" }
}),
is_enabled: false,
created_at_unix_secs: 1_710_000_000,
created_at_unix_ms: 1_710_000_000,
updated_at_unix_secs: 1_710_000_050,
},
]),
@@ -516,7 +516,7 @@ async fn gateway_handles_admin_billing_collector_routes_locally_with_trusted_adm
default_value: None,
priority: 10,
is_enabled: true,
created_at_unix_secs: 1_710_000_000,
created_at_unix_ms: 1_710_000_000,
updated_at_unix_secs: 1_710_000_100,
},
crate::AdminBillingCollectorRecord {
@@ -531,7 +531,7 @@ async fn gateway_handles_admin_billing_collector_routes_locally_with_trusted_adm
default_value: Some("1".to_string()),
priority: 5,
is_enabled: false,
created_at_unix_secs: 1_710_000_000,
created_at_unix_ms: 1_710_000_000,
updated_at_unix_secs: 1_710_000_050,
},
]),

View File

@@ -46,7 +46,7 @@ async fn gateway_handles_admin_provider_keys_locally_with_trusted_admin_principa
key_a.success_count = Some(9);
key_a.error_count = Some(3);
key_a.total_response_time_ms = Some(1800);
key_a.created_at_unix_secs = Some(1_711_000_000);
key_a.created_at_unix_ms = Some(1_711_000_000);
key_a.updated_at_unix_secs = Some(1_711_000_100);
key_a.last_used_at_unix_secs = Some(1_711_000_120);
key_a.note = Some("primary key".to_string());
@@ -67,7 +67,7 @@ async fn gateway_handles_admin_provider_keys_locally_with_trusted_admin_principa
key_b.success_count = Some(4);
key_b.error_count = Some(0);
key_b.total_response_time_ms = Some(400);
key_b.created_at_unix_secs = Some(1_711_100_000);
key_b.created_at_unix_ms = Some(1_711_100_000);
key_b.updated_at_unix_secs = Some(1_711_100_100);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
@@ -895,7 +895,7 @@ async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin
key_a.internal_priority = 10;
key_a.request_count = Some(12);
key_a.success_count = Some(9);
key_a.created_at_unix_secs = Some(1_711_000_000);
key_a.created_at_unix_ms = Some(1_711_000_000);
key_a.updated_at_unix_secs = Some(1_711_000_100);
key_a.capabilities = Some(json!({"cache_1h": true, "gemini_files": false}));
key_a.global_priority_by_format = Some(json!({"openai:chat": 3}));
@@ -906,7 +906,7 @@ async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin
key_b.internal_priority = 20;
key_b.request_count = Some(2);
key_b.success_count = Some(1);
key_b.created_at_unix_secs = Some(1_711_100_000);
key_b.created_at_unix_ms = Some(1_711_100_000);
key_b.updated_at_unix_secs = Some(1_711_100_100);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(

View File

@@ -124,7 +124,7 @@ fn sample_usage(
total_cost_usd: f64,
status: &str,
status_code: Option<i32>,
created_at_unix_secs: i64,
created_at_unix_ms: i64,
) -> StoredRequestUsageAudit {
let is_error = status_code.is_some_and(|value| value >= 400)
|| status.trim().eq_ignore_ascii_case("failed")
@@ -163,9 +163,9 @@ fn sample_usage(
Some(30),
status.to_string(),
"billed".to_string(),
created_at_unix_secs,
created_at_unix_secs,
Some(created_at_unix_secs),
created_at_unix_ms,
created_at_unix_ms,
Some(created_at_unix_ms),
)
.expect("usage should build")
}
@@ -285,7 +285,7 @@ fn sample_candidate(
request_id: &str,
candidate_index: i32,
status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>,
started_at_unix_ms: Option<i64>,
latency_ms: Option<i32>,
status_code: Option<i32>,
) -> StoredRequestCandidate {
@@ -311,9 +311,9 @@ fn sample_candidate(
Some(1),
None,
Some(json!({"cache_1h": true})),
100 + i64::from(candidate_index),
started_at_unix_secs,
started_at_unix_secs.map(|value| value + 1),
(100 + i64::from(candidate_index)) * 1_000,
started_at_unix_ms.map(|v| v * 1_000),
started_at_unix_ms.map(|value| (value + 1) * 1_000),
)
.expect("candidate should build")
}

View File

@@ -338,7 +338,7 @@ async fn gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_a
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": "proxy-node-kiro",
"created_at_unix_secs": 1_711_000_000u64,
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
@@ -473,7 +473,7 @@ async fn local_admin_provider_oauth_device_poll_attaches_audit_only_when_transit
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": null,
"created_at_unix_secs": 1_711_000_000u64,
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
@@ -522,7 +522,7 @@ async fn local_admin_provider_oauth_device_poll_attaches_audit_only_when_transit
"expires_at_unix_secs": 4_102_444_800u64,
"status": "authorized",
"proxy_node_id": null,
"created_at_unix_secs": 1_711_000_000u64,
"created_at_unix_ms": 1_711_000_000u64,
"key_id": "key-existing",
"email": "kiro@example.com",
"replaced": false,

View File

@@ -68,7 +68,7 @@ fn sample_payment_order(
amount_usd: f64,
payment_method: &str,
status: &str,
created_at_unix_secs: u64,
created_at_unix_ms: u64,
expires_at_unix_secs: Option<u64>,
) -> crate::AdminWalletPaymentOrderRecord {
crate::AdminWalletPaymentOrderRecord {
@@ -86,7 +86,7 @@ fn sample_payment_order(
gateway_order_id: None,
status: status.to_string(),
gateway_response: None,
created_at_unix_secs,
created_at_unix_ms,
paid_at_unix_secs: None,
credited_at_unix_secs: None,
expires_at_unix_secs,
@@ -99,7 +99,7 @@ fn sample_payment_callback(
payment_method: &str,
callback_key: &str,
status: &str,
created_at_unix_secs: u64,
created_at_unix_ms: u64,
) -> crate::state::AdminPaymentCallbackRecord {
crate::state::AdminPaymentCallbackRecord {
id: callback_id.to_string(),
@@ -113,8 +113,8 @@ fn sample_payment_callback(
status: status.to_string(),
payload: Some(json!({ "source": callback_id })),
error_message: (status == "failed").then(|| "signature mismatch".to_string()),
created_at_unix_secs,
processed_at_unix_secs: Some(created_at_unix_secs.saturating_add(60)),
created_at_unix_ms,
processed_at_unix_secs: Some(created_at_unix_ms.saturating_add(60)),
}
}

View File

@@ -1074,6 +1074,9 @@ async fn gateway_handles_admin_provider_health_monitor_locally_with_trusted_admi
assert_eq!(endpoints[0]["failed_count"], 1);
assert_eq!(endpoints[0]["skipped_count"], 0);
assert_eq!(endpoints[0]["success_rate"], 0.5);
assert!(endpoints[0]["last_event_at"]
.as_str()
.is_some_and(|value| value.ends_with(".000Z")));
assert_eq!(
endpoints[0]["events"]
.as_array()
@@ -1081,6 +1084,9 @@ async fn gateway_handles_admin_provider_health_monitor_locally_with_trusted_admi
.len(),
2
);
assert!(endpoints[0]["events"][0]["timestamp"]
.as_str()
.is_some_and(|value| value.ends_with(".000Z")));
assert_eq!(endpoints[1]["endpoint_id"], "endpoint-openai-cli");
assert_eq!(endpoints[1]["api_format"], "openai:cli");
assert_eq!(endpoints[1]["total_attempts"], 1);

View File

@@ -167,14 +167,14 @@ async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_prin
node_id: "node-1".to_string(),
event_type: "connected".to_string(),
detail: Some("older".to_string()),
created_at_unix_secs: Some(1_710_000_000),
created_at_unix_ms: Some(1_710_000_000),
},
StoredProxyNodeEvent {
id: 2,
node_id: "node-1".to_string(),
event_type: "disconnected".to_string(),
detail: Some("newer".to_string()),
created_at_unix_secs: Some(1_710_000_100),
created_at_unix_ms: Some(1_710_000_100),
},
],
));

View File

@@ -68,7 +68,7 @@ fn sample_usage_row(
output_tokens: i32,
total_cost_usd: f64,
actual_total_cost_usd: f64,
created_at_unix_secs: i64,
created_at_unix_ms: i64,
) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
id.to_string(),
@@ -104,9 +104,9 @@ fn sample_usage_row(
Some(120),
"completed".to_string(),
"settled".to_string(),
created_at_unix_secs,
created_at_unix_secs + 1,
Some(created_at_unix_secs + 2),
created_at_unix_ms,
created_at_unix_ms + 1,
Some(created_at_unix_ms + 2),
)
.expect("usage row should build")
}

View File

@@ -149,7 +149,7 @@ fn sample_usage_row(
output_tokens: i32,
total_cost_usd: f64,
actual_total_cost_usd: f64,
created_at_unix_secs: i64,
created_at_unix_ms: i64,
) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
id.to_string(),
@@ -185,9 +185,9 @@ fn sample_usage_row(
Some(120),
status.to_string(),
"settled".to_string(),
created_at_unix_secs,
created_at_unix_secs + 1,
Some(created_at_unix_secs + 2),
created_at_unix_ms,
created_at_unix_ms + 1,
Some(created_at_unix_ms + 2),
)
.expect("usage row should build")
.with_cache_input_tokens(15, 5)

View File

@@ -79,7 +79,7 @@ async fn local_admin_video_tasks_response(
fn sample_admin_video_task(
id: &str,
status: VideoTaskStatus,
created_at_unix_secs: u64,
created_at_unix_ms: u64,
user_id: &str,
username: &str,
provider_id: &str,
@@ -120,14 +120,14 @@ fn sample_admin_video_task(
next_poll_at_unix_secs: None,
poll_count: 1,
max_poll_count: 360,
created_at_unix_secs,
submitted_at_unix_secs: Some(created_at_unix_secs),
created_at_unix_ms,
submitted_at_unix_secs: Some(created_at_unix_ms),
completed_at_unix_secs: if matches!(status, VideoTaskStatus::Completed) {
Some(created_at_unix_secs + 30)
Some(created_at_unix_ms + 30)
} else {
None
},
updated_at_unix_secs: created_at_unix_secs + 5,
updated_at_unix_secs: created_at_unix_ms + 5,
error_code: None,
error_message: None,
video_url: Some(format!("https://example.com/{id}.mp4")),
@@ -542,7 +542,7 @@ async fn gateway_cancels_admin_video_task_locally_with_trusted_admin_principal()
"OpenAi": {
"local_task_id": "task-openai-cancel",
"upstream_task_id": "ext-video-task-123",
"created_at_unix_secs": 1710000400,
"created_at_unix_ms": 1710000400,
"user_id": "user-1",
"api_key_id": "api-key-task-openai-cancel",
"model": "sora-2",

View File

@@ -146,7 +146,7 @@ fn sample_payment_order_record(
gateway_order_id: None,
status: "credited".to_string(),
gateway_response: None,
created_at_unix_secs: 1_710_000_000,
created_at_unix_ms: 1_710_000_000,
paid_at_unix_secs: Some(1_710_000_060),
credited_at_unix_secs: Some(1_710_000_120),
expires_at_unix_secs: None,
@@ -174,7 +174,7 @@ fn sample_transaction_record(
link_id: Some("order-123".to_string()),
operator_id: operator_id.map(ToOwned::to_owned),
description: Some("Admin manual credit".to_string()),
created_at_unix_secs: 1_710_000_200,
created_at_unix_ms: 1_710_000_200,
}
}
@@ -210,7 +210,7 @@ fn sample_refund_record(
requested_by: user_id.map(ToOwned::to_owned),
approved_by: approved_by.map(ToOwned::to_owned),
processed_by: processed_by.map(ToOwned::to_owned),
created_at_unix_secs: 1_710_000_000,
created_at_unix_ms: 1_710_000_000,
updated_at_unix_secs: processed_at_unix_secs.unwrap_or(1_710_000_000),
processed_at_unix_secs,
completed_at_unix_secs: None,
@@ -432,7 +432,7 @@ async fn gateway_handles_admin_wallets_refund_requests_locally_with_trusted_admi
None,
None,
);
refund.created_at_unix_secs = 1_710_000_300;
refund.created_at_unix_ms = 1_710_000_300;
refund
},
{
@@ -447,7 +447,7 @@ async fn gateway_handles_admin_wallets_refund_requests_locally_with_trusted_admi
None,
None,
);
refund.created_at_unix_secs = 1_710_000_200;
refund.created_at_unix_ms = 1_710_000_200;
refund
},
sample_refund_record(
@@ -612,7 +612,7 @@ async fn gateway_handles_admin_wallets_detail_locally_with_trusted_admin_princip
"wallet-user-123",
Some("admin-user-operator"),
);
transaction.created_at_unix_secs = 1_710_000_300;
transaction.created_at_unix_ms = 1_710_000_300;
transaction.description = Some("Most recent credit".to_string());
transaction
},
@@ -803,7 +803,7 @@ async fn gateway_handles_admin_wallets_transactions_locally_with_trusted_admin_p
"wallet-user-123",
Some("admin-user-operator"),
);
transaction.created_at_unix_secs = 1_710_000_300;
transaction.created_at_unix_ms = 1_710_000_300;
transaction.description = Some("Most recent credit".to_string());
transaction
},
@@ -958,7 +958,7 @@ async fn gateway_handles_admin_wallets_refunds_locally_with_trusted_admin_princi
None,
None,
);
refund.created_at_unix_secs = 1_710_000_300;
refund.created_at_unix_ms = 1_710_000_300;
refund
},
{
@@ -973,7 +973,7 @@ async fn gateway_handles_admin_wallets_refunds_locally_with_trusted_admin_princi
Some("admin-user-123"),
Some(1_710_000_250),
);
refund.created_at_unix_secs = 1_710_000_200;
refund.created_at_unix_ms = 1_710_000_200;
refund
},
sample_refund_record(

View File

@@ -503,6 +503,8 @@ pub(super) fn sample_request_candidate(
created_at_unix_secs: i64,
finished_at_unix_secs: Option<i64>,
) -> StoredRequestCandidate {
let created_at_unix_ms = created_at_unix_secs * 1_000;
let finished_at_unix_ms = finished_at_unix_secs.map(|v| v * 1_000);
StoredRequestCandidate::new(
id.to_string(),
request_id.to_string(),
@@ -525,9 +527,9 @@ pub(super) fn sample_request_candidate(
Some(1),
None,
None,
created_at_unix_secs,
Some(created_at_unix_secs),
finished_at_unix_secs,
created_at_unix_ms,
Some(created_at_unix_ms),
finished_at_unix_ms,
)
.expect("request candidate should build")
}
@@ -562,9 +564,9 @@ pub(super) fn sample_recent_key_rpm_candidate(
Some(concurrent_requests),
None,
None,
now_unix_secs - 10,
Some(now_unix_secs - 10),
Some(now_unix_secs - 8),
(now_unix_secs - 10) * 1_000,
Some((now_unix_secs - 10) * 1_000),
Some((now_unix_secs - 8) * 1_000),
)
.expect("request candidate should build")
}