mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
refactor: 抽离 AI pipeline 与调度共享能力逻辑
This commit is contained in:
@@ -25,461 +25,6 @@ use aether_data_contracts::repository::provider_catalog::{
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_compact_cross_format_upstream_stream_via_local_finalize_response()
|
||||
{
|
||||
use base64::Engine as _;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenRemoteExecutionRuntimeRequest {
|
||||
trace_id: String,
|
||||
request_id: String,
|
||||
url: String,
|
||||
authorization: String,
|
||||
endpoint_tag: String,
|
||||
provider_model: String,
|
||||
has_model_field: bool,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai", "gemini"])),
|
||||
Some(serde_json::json!(["openai:compact"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800_i64),
|
||||
Some(serde_json::json!(["openai", "gemini"])),
|
||||
Some(serde_json::json!(["openai:compact"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-openai-compact-gemini-finalize-local-1".to_string(),
|
||||
provider_name: "gemini".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-openai-compact-gemini-finalize-local-1".to_string(),
|
||||
endpoint_api_format: "gemini:cli".to_string(),
|
||||
endpoint_api_family: Some("gemini".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-openai-compact-gemini-finalize-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "bearer".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["gemini:cli".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"gemini:cli": 1})),
|
||||
model_id: "model-openai-compact-gemini-finalize-local-1".to_string(),
|
||||
global_model_id: "global-model-openai-compact-gemini-finalize-local-1".to_string(),
|
||||
global_model_name: "gpt-5".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gemini-2.5-pro-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gemini-2.5-pro-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["gemini:cli".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-openai-compact-gemini-finalize-local-1".to_string(),
|
||||
"gemini".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-openai-compact-gemini-finalize-local-1".to_string(),
|
||||
"provider-openai-compact-gemini-finalize-local-1".to_string(),
|
||||
"gemini:cli".to_string(),
|
||||
Some("gemini".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://generativelanguage.googleapis.com".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"openai-compact-gemini-finalize-cross-format"}
|
||||
])),
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-openai-compact-gemini-finalize-local-1".to_string(),
|
||||
"provider-openai-compact-gemini-finalize-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["gemini:cli"])),
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"sk-upstream-openai-compact-gemini-finalize",
|
||||
)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"gemini:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_remote_execution_runtime =
|
||||
Arc::new(Mutex::new(None::<SeenRemoteExecutionRuntimeRequest>));
|
||||
let seen_remote_execution_runtime_clone = Arc::clone(&seen_remote_execution_runtime);
|
||||
let report_hits = Arc::new(Mutex::new(0usize));
|
||||
let report_hits_clone = Arc::clone(&report_hits);
|
||||
let finalize_hits = Arc::new(Mutex::new(0usize));
|
||||
let finalize_hits_clone = Arc::clone(&finalize_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 request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
|
||||
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/finalize-sync",
|
||||
any(move |_request: Request| {
|
||||
let finalize_hits_inner = Arc::clone(&finalize_hits_clone);
|
||||
async move {
|
||||
*finalize_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(
|
||||
StatusCode::IM_A_TEAPOT,
|
||||
Body::from("finalize-sync-should-not-be-hit"),
|
||||
)
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |request: Request| {
|
||||
let report_hits_inner = Arc::clone(&report_hits_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*report_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/responses/compact",
|
||||
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(move |request: Request| {
|
||||
let seen_remote_execution_runtime_inner = Arc::clone(&seen_remote_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
|
||||
*seen_remote_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenRemoteExecutionRuntimeRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
request_id: payload
|
||||
.get("request_id")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
authorization: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("authorization"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
provider_model: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
has_model_field: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.is_some(),
|
||||
});
|
||||
Json(json!({
|
||||
"request_id": "trace-openai-compact-xfmt-stream-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "text/event-stream"
|
||||
},
|
||||
"body": {
|
||||
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
|
||||
concat!(
|
||||
"data: {\"responseId\":\"upstream-compact-stream-123\",\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hello \"}],\"role\":\"model\"},\"index\":0}],\"modelVersion\":\"gemini-2.5-pro-upstream\"}\n\n",
|
||||
"data: {\"responseId\":\"upstream-compact-stream-123\",\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hello Gemini Compact\"}],\"role\":\"model\"},\"finishReason\":\"STOP\",\"index\":0}],\"modelVersion\":\"gemini-2.5-pro-upstream\",\"usageMetadata\":{\"promptTokenCount\":2,\"candidatesTokenCount\":3,\"totalTokenCount\":5}}\n\n"
|
||||
)
|
||||
)
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 31
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-compact-xfmt-stream")),
|
||||
sample_auth_snapshot(
|
||||
"api-key-openai-compact-xfmt-stream-1",
|
||||
"user-openai-compact-xfmt-stream-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.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,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let started_at = std::time::Instant::now();
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/responses/compact"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-compact-xfmt-stream",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-openai-compact-xfmt-stream-123")
|
||||
.body("{\"model\":\"gpt-5\",\"input\":\"hello\"}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
let elapsed = started_at.elapsed();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let response_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
response_json,
|
||||
json!({
|
||||
"id": "upstream-compact-stream-123",
|
||||
"object": "response",
|
||||
"status": "completed",
|
||||
"model": "gemini-2.5-pro-upstream",
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": "upstream-compact-stream-123_msg",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{
|
||||
"type": "output_text",
|
||||
"text": "Hello Gemini Compact",
|
||||
"annotations": []
|
||||
}]
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 2,
|
||||
"output_tokens": 3,
|
||||
"total_tokens": 5
|
||||
}
|
||||
})
|
||||
);
|
||||
assert!(
|
||||
elapsed < std::time::Duration::from_millis(10_000),
|
||||
"response took unexpectedly long for local finalize path: elapsed={elapsed:?} finalize_hits={} report_hits={}",
|
||||
*finalize_hits.lock().expect("mutex should lock"),
|
||||
*report_hits.lock().expect("mutex should lock"),
|
||||
);
|
||||
|
||||
let seen_remote_execution_runtime_request = seen_remote_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("remote execution runtime plan should be captured");
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.trace_id,
|
||||
"trace-openai-compact-xfmt-stream-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.request_id,
|
||||
"trace-openai-compact-xfmt-stream-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.url,
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro-upstream:generateContent"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.authorization,
|
||||
"Bearer sk-upstream-openai-compact-gemini-finalize"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.endpoint_tag,
|
||||
"openai-compact-gemini-finalize-cross-format"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.provider_model,
|
||||
"gemini-2.5-pro-upstream"
|
||||
);
|
||||
assert!(seen_remote_execution_runtime_request.has_model_field);
|
||||
|
||||
let mut stored_candidates = Vec::new();
|
||||
for _ in 0..50 {
|
||||
stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-openai-compact-xfmt-stream-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
if stored_candidates.len() == 1
|
||||
&& stored_candidates[0].status == RequestCandidateStatus::Success
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert_eq!(
|
||||
*report_hits.lock().expect("mutex should lock"),
|
||||
0,
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
assert_eq!(*finalize_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_executes_openai_compact_openai_family_upstream_stream_via_local_finalize_response()
|
||||
{
|
||||
@@ -716,12 +261,13 @@ async fn gateway_executes_openai_compact_openai_family_upstream_stream_via_local
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_remote_execution_runtime_inner = Arc::clone(&seen_remote_execution_runtime_clone);
|
||||
let seen_remote_execution_runtime_inner =
|
||||
Arc::clone(&seen_remote_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
|
||||
.expect("execution runtime payload should parse");
|
||||
*seen_remote_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenRemoteExecutionRuntimeRequest {
|
||||
@@ -808,18 +354,17 @@ async fn gateway_executes_openai_compact_openai_family_upstream_stream_via_local
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
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,
|
||||
),
|
||||
);
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.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,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
@@ -926,451 +471,3 @@ async fn gateway_executes_openai_compact_openai_family_upstream_stream_via_local
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_compact_openai_family_upstream_stream_via_local_finalize_response_even_when_conversion_flagged(
|
||||
) {
|
||||
use base64::Engine as _;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenRemoteExecutionRuntimeRequest {
|
||||
trace_id: String,
|
||||
request_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
authorization: String,
|
||||
endpoint_tag: String,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:compact"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800_i64),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:compact"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-openai-compact-openai-family-conversion-local-1".to_string(),
|
||||
provider_name: "openai".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-openai-compact-openai-family-conversion-local-1".to_string(),
|
||||
endpoint_api_format: "openai:cli".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-openai-compact-openai-family-conversion-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "bearer".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:cli".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:cli": 1})),
|
||||
model_id: "model-openai-compact-openai-family-conversion-local-1".to_string(),
|
||||
global_model_id: "global-model-openai-compact-openai-family-conversion-local-1"
|
||||
.to_string(),
|
||||
global_model_name: "gpt-5".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-5-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:cli".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-openai-compact-openai-family-conversion-local-1".to_string(),
|
||||
"openai".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-openai-compact-openai-family-conversion-local-1".to_string(),
|
||||
"provider-openai-compact-openai-family-conversion-local-1".to_string(),
|
||||
"openai:cli".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://api.openai.example".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"openai-cli-compact-family-finalize-local"}
|
||||
])),
|
||||
None,
|
||||
Some(2),
|
||||
Some("/custom/v1/responses".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-openai-compact-openai-family-conversion-local-1".to_string(),
|
||||
"provider-openai-compact-openai-family-conversion-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["openai:cli"])),
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"sk-upstream-openai-cli-family-finalize",
|
||||
)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"openai:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_remote_execution_runtime =
|
||||
Arc::new(Mutex::new(None::<SeenRemoteExecutionRuntimeRequest>));
|
||||
let seen_remote_execution_runtime_clone = Arc::clone(&seen_remote_execution_runtime);
|
||||
let report_hits = Arc::new(Mutex::new(0usize));
|
||||
let report_hits_clone = Arc::clone(&report_hits);
|
||||
let finalize_hits = Arc::new(Mutex::new(0usize));
|
||||
let finalize_hits_clone = Arc::clone(&finalize_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 request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
|
||||
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/finalize-sync",
|
||||
any(move |_request: Request| {
|
||||
let finalize_hits_inner = Arc::clone(&finalize_hits_clone);
|
||||
async move {
|
||||
*finalize_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(
|
||||
StatusCode::IM_A_TEAPOT,
|
||||
Body::from("finalize-sync-should-not-be-hit"),
|
||||
)
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |request: Request| {
|
||||
let report_hits_inner = Arc::clone(&report_hits_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*report_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/responses/compact",
|
||||
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(move |request: Request| {
|
||||
let seen_remote_execution_runtime_inner = Arc::clone(&seen_remote_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
|
||||
*seen_remote_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenRemoteExecutionRuntimeRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
request_id: payload
|
||||
.get("request_id")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
model: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
authorization: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("authorization"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
Json(json!({
|
||||
"request_id": "trace-openai-compact-openai-family-conversion-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "text/event-stream"
|
||||
},
|
||||
"body": {
|
||||
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
|
||||
concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_compact_openai_family_conversion_123\",\"object\":\"response\",\"model\":\"gpt-5\",\"status\":\"in_progress\",\"output\":[]}}\n\n",
|
||||
"event: response.output_text.delta\n",
|
||||
"data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"Hello Compact\"}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_compact_openai_family_conversion_123\",\"object\":\"response\",\"model\":\"gpt-5\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"total_tokens\":5}}}\n\n"
|
||||
)
|
||||
)
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 31
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key(
|
||||
"sk-client-openai-compact-openai-family-conversion",
|
||||
)),
|
||||
sample_auth_snapshot(
|
||||
"api-key-openai-compact-openai-family-conversion-1",
|
||||
"user-openai-compact-openai-family-conversion-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.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,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let started_at = std::time::Instant::now();
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/responses/compact"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-compact-openai-family-conversion",
|
||||
)
|
||||
.header(
|
||||
TRACE_ID_HEADER,
|
||||
"trace-openai-compact-openai-family-conversion-123",
|
||||
)
|
||||
.body("{\"model\":\"gpt-5\",\"input\":\"hello\"}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
let elapsed = started_at.elapsed();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let response_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
response_json,
|
||||
json!({
|
||||
"id": "resp_compact_openai_family_conversion_123",
|
||||
"object": "response",
|
||||
"model": "gpt-5",
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"usage": {
|
||||
"input_tokens": 2,
|
||||
"output_tokens": 3,
|
||||
"total_tokens": 5
|
||||
}
|
||||
})
|
||||
);
|
||||
assert!(
|
||||
elapsed < std::time::Duration::from_millis(10_000),
|
||||
"response took unexpectedly long for local finalize path: elapsed={elapsed:?} finalize_hits={} report_hits={}",
|
||||
*finalize_hits.lock().expect("mutex should lock"),
|
||||
*report_hits.lock().expect("mutex should lock"),
|
||||
);
|
||||
|
||||
let seen_remote_execution_runtime_request = seen_remote_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("remote execution runtime plan should be captured");
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.trace_id,
|
||||
"trace-openai-compact-openai-family-conversion-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.request_id,
|
||||
"trace-openai-compact-openai-family-conversion-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.url,
|
||||
"https://api.openai.example/custom/v1/responses"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.model,
|
||||
"gpt-5-upstream"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.authorization,
|
||||
"Bearer sk-upstream-openai-cli-family-finalize"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_remote_execution_runtime_request.endpoint_tag,
|
||||
"openai-cli-compact-family-finalize-local"
|
||||
);
|
||||
|
||||
let mut stored_candidates = Vec::new();
|
||||
for _ in 0..50 {
|
||||
stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-openai-compact-openai-family-conversion-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
if stored_candidates.len() == 1
|
||||
&& stored_candidates[0].status == RequestCandidateStatus::Success
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert_eq!(
|
||||
*report_hits.lock().expect("mutex should lock"),
|
||||
0,
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
assert_eq!(*finalize_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();
|
||||
}
|
||||
|
||||
@@ -1895,8 +1895,8 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_with_local_fail
|
||||
assert_eq!(stored_candidates[1].candidate_index, 1);
|
||||
assert_eq!(stored_candidates[1].status, RequestCandidateStatus::Success);
|
||||
assert_eq!(stored_candidates[1].status_code, Some(200));
|
||||
assert!(stored_candidates[1].started_at_unix_secs.is_some());
|
||||
assert!(stored_candidates[1].finished_at_unix_secs.is_some());
|
||||
assert!(stored_candidates[1].started_at_unix_ms.is_some());
|
||||
assert!(stored_candidates[1].finished_at_unix_ms.is_some());
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
|
||||
@@ -396,8 +396,8 @@ async fn gateway_skips_unsupported_local_openai_chat_sync_candidate_before_tryin
|
||||
skipped_candidate.skip_reason.as_deref(),
|
||||
Some("transport_unsupported")
|
||||
);
|
||||
assert!(skipped_candidate.started_at_unix_secs.is_none());
|
||||
assert!(skipped_candidate.finished_at_unix_secs.is_some());
|
||||
assert!(skipped_candidate.started_at_unix_ms.is_none());
|
||||
assert!(skipped_candidate.finished_at_unix_ms.is_some());
|
||||
let successful_candidate = stored_candidates
|
||||
.iter()
|
||||
.find(|candidate| candidate.candidate_index == 1)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -6,7 +6,8 @@ use super::{
|
||||
Request, RequestCandidateReadRepository, RequestCandidateStatus, Router, Sha256, StatusCode,
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER,
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -425,3 +426,313 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_claude_chat_error_for_local_sync_failure() {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:chat"])),
|
||||
Some(serde_json::json!(["claude-sonnet-4-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:chat"])),
|
||||
Some(serde_json::json!(["claude-sonnet-4-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-claude-local-1".to_string(),
|
||||
provider_name: "claude".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-claude-local-1".to_string(),
|
||||
endpoint_api_format: "claude:chat".to_string(),
|
||||
endpoint_api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-claude-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["claude:chat".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"claude:chat": 1})),
|
||||
model_id: "model-claude-local-1".to_string(),
|
||||
global_model_id: "global-model-claude-local-1".to_string(),
|
||||
global_model_name: "claude-sonnet-4-5".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "claude-sonnet-4-5-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "claude-sonnet-4-5-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["claude:chat".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-claude-local-1".to_string(),
|
||||
"claude".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-claude-local-1".to_string(),
|
||||
"provider-claude-local-1".to_string(),
|
||||
"claude:chat".to_string(),
|
||||
Some("claude".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://api.anthropic.example".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"claude-chat-local"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"metadata.mode","value":"safe"},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"}
|
||||
])),
|
||||
Some(2),
|
||||
Some("/custom/v1/messages".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-claude-local-1".to_string(),
|
||||
"provider-claude-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["claude:chat"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-claude-chat")
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"claude:chat": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-claude-chat-local"})),
|
||||
Some(serde_json::json!({"tls_profile":"chrome_136"})),
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
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 seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/messages",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |_request: Request| async move {
|
||||
Json(json!({
|
||||
"request_id": "trace-claude-chat-local-error-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "rate_limit_error",
|
||||
"message": "slow down"
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 29
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-claude-chat-local-error")),
|
||||
sample_auth_snapshot("api-key-claude-local-error-1", "user-claude-local-error-1"),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("x-api-key", "sk-client-claude-chat-local-error")
|
||||
.header(TRACE_ID_HEADER, "trace-claude-chat-local-error-123")
|
||||
.body(
|
||||
"{\"model\":\"claude-sonnet-4-5\",\"messages\":[],\"metadata\":{\"client\":\"desktop-claude\"}}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
|
||||
);
|
||||
let response_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
response_json,
|
||||
json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "rate_limit_error",
|
||||
"message": "slow down"
|
||||
}
|
||||
})
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-claude-chat-local-error-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
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();
|
||||
}
|
||||
|
||||
@@ -6,7 +6,8 @@ use super::{
|
||||
Request, RequestCandidateReadRepository, RequestCandidateStatus, Router, Sha256, StatusCode,
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER,
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -433,3 +434,277 @@ async fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_syn
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_claude_cli_error_for_local_sync_failure() {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["claude-code"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["claude-code"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-claude-cli-local-1".to_string(),
|
||||
provider_name: "claude".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-claude-cli-local-1".to_string(),
|
||||
endpoint_api_format: "claude:cli".to_string(),
|
||||
endpoint_api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-claude-cli-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "bearer".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["claude:cli".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"claude:cli": 1})),
|
||||
model_id: "model-claude-cli-local-1".to_string(),
|
||||
global_model_id: "global-model-claude-cli-local-1".to_string(),
|
||||
global_model_name: "claude-code".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "claude-code-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "claude-code-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["claude:cli".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-claude-cli-local-1".to_string(),
|
||||
"claude".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-claude-cli-local-1".to_string(),
|
||||
"provider-claude-cli-local-1".to_string(),
|
||||
"claude:cli".to_string(),
|
||||
Some("claude".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://api.anthropic.example".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"claude-cli-local"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"metadata.mode","value":"safe"},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"}
|
||||
])),
|
||||
Some(2),
|
||||
Some("/custom/v1/messages".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-claude-cli-local-1".to_string(),
|
||||
"provider-claude-cli-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-claude-cli")
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"claude:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-claude-cli-local"})),
|
||||
Some(serde_json::json!({"tls_profile":"chrome_136"})),
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |_request: Request| async move {
|
||||
Json(json!({
|
||||
"request_id": "trace-claude-cli-local-error-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "rate_limit_error",
|
||||
"message": "slow down"
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 29
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-claude-cli-local-error")),
|
||||
sample_auth_snapshot(
|
||||
"api-key-claude-cli-local-error-1",
|
||||
"user-claude-cli-local-error-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-claude-cli-local-error",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-claude-cli-local-error-123")
|
||||
.body(
|
||||
"{\"model\":\"claude-code\",\"messages\":[],\"metadata\":{\"client\":\"desktop-claude-cli\"}}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
|
||||
);
|
||||
let response_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
response_json,
|
||||
json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "rate_limit_error",
|
||||
"message": "slow down"
|
||||
}
|
||||
})
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-claude-cli-local-error-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -6,7 +6,8 @@ use super::{
|
||||
Request, RequestCandidateReadRepository, RequestCandidateStatus, Router, Sha256, StatusCode,
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER,
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -444,6 +445,277 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_with_local_syn
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_gemini_cli_error_for_local_sync_failure() {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["gemini"])),
|
||||
Some(serde_json::json!(["gemini:cli"])),
|
||||
Some(serde_json::json!(["gemini-cli"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["gemini"])),
|
||||
Some(serde_json::json!(["gemini:cli"])),
|
||||
Some(serde_json::json!(["gemini-cli"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-gemini-cli-local-1".to_string(),
|
||||
provider_name: "gemini".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-gemini-cli-local-1".to_string(),
|
||||
endpoint_api_format: "gemini:cli".to_string(),
|
||||
endpoint_api_family: Some("gemini".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-gemini-cli-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "bearer".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["gemini:cli".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"gemini:cli": 1})),
|
||||
model_id: "model-gemini-cli-local-1".to_string(),
|
||||
global_model_id: "global-model-gemini-cli-local-1".to_string(),
|
||||
global_model_name: "gemini-cli".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gemini-cli-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gemini-cli-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["gemini:cli".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-gemini-cli-local-1".to_string(),
|
||||
"gemini".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-gemini-cli-local-1".to_string(),
|
||||
"provider-gemini-cli-local-1".to_string(),
|
||||
"gemini:cli".to_string(),
|
||||
Some("gemini".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://generativelanguage.googleapis.com".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"gemini-cli-local"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"metadata.mode","value":"safe"},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"},
|
||||
{"action":"drop","path":"toolConfig"}
|
||||
])),
|
||||
Some(2),
|
||||
Some("/custom/v1beta/models/gemini-cli-upstream:generateContent".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-gemini-cli-local-1".to_string(),
|
||||
"provider-gemini-cli-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["gemini:cli"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-gemini-cli")
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"gemini:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-gemini-cli-local"})),
|
||||
Some(serde_json::json!({"tls_profile":"chrome_136"})),
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |_request: Request| async move {
|
||||
Json(json!({
|
||||
"request_id": "trace-gemini-cli-local-error-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"error": {
|
||||
"message": "quota reached",
|
||||
"status": "RESOURCE_EXHAUSTED"
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 27
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("client-gemini-cli-local-error")),
|
||||
sample_auth_snapshot(
|
||||
"api-key-gemini-cli-local-error-1",
|
||||
"user-gemini-cli-local-error-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1beta/models/gemini-cli:generateContent"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("user-agent", "GeminiCLI/1.0")
|
||||
.header("x-goog-api-key", "client-gemini-cli-local-error")
|
||||
.header(TRACE_ID_HEADER, "trace-gemini-cli-local-error-123")
|
||||
.body(
|
||||
"{\"contents\":[],\"generationConfig\":{\"temperature\":0.2},\"metadata\":{\"client\":\"desktop-gemini-cli\"},\"toolConfig\":{\"functionCallingConfig\":{\"mode\":\"AUTO\"}}}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
|
||||
);
|
||||
let response_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
response_json,
|
||||
json!({
|
||||
"error": {
|
||||
"message": "quota reached",
|
||||
"status": "RESOURCE_EXHAUSTED"
|
||||
}
|
||||
})
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-gemini-cli-local-error-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_refresh() {
|
||||
#[derive(Debug, Clone)]
|
||||
|
||||
@@ -6,7 +6,8 @@ use super::{
|
||||
Request, RequestCandidateReadRepository, RequestCandidateStatus, Router, Sha256, StatusCode,
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER,
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -433,3 +434,313 @@ async fn gateway_executes_gemini_chat_sync_via_local_decision_gate_with_local_sy
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_gemini_chat_error_for_local_sync_failure() {
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["gemini"])),
|
||||
Some(serde_json::json!(["gemini:chat"])),
|
||||
Some(serde_json::json!(["gemini-2.5-pro"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["gemini"])),
|
||||
Some(serde_json::json!(["gemini:chat"])),
|
||||
Some(serde_json::json!(["gemini-2.5-pro"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-gemini-local-1".to_string(),
|
||||
provider_name: "gemini".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-gemini-local-1".to_string(),
|
||||
endpoint_api_format: "gemini:chat".to_string(),
|
||||
endpoint_api_family: Some("gemini".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-gemini-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["gemini:chat".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"gemini:chat": 1})),
|
||||
model_id: "model-gemini-local-1".to_string(),
|
||||
global_model_id: "global-model-gemini-local-1".to_string(),
|
||||
global_model_name: "gemini-2.5-pro".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gemini-2.5-pro-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gemini-2.5-pro-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["gemini:chat".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-gemini-local-1".to_string(),
|
||||
"gemini".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-gemini-local-1".to_string(),
|
||||
"provider-gemini-local-1".to_string(),
|
||||
"gemini:chat".to_string(),
|
||||
Some("gemini".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://generativelanguage.googleapis.com".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"gemini-chat-local"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"metadata.mode","value":"safe"},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"},
|
||||
{"action":"drop","path":"toolConfig"}
|
||||
])),
|
||||
Some(2),
|
||||
Some("/custom/v1beta/models/gemini-2.5-pro-upstream:generateContent".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-gemini-local-1".to_string(),
|
||||
"provider-gemini-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["gemini:chat"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-gemini-chat")
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"gemini:chat": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-gemini-chat-local"})),
|
||||
Some(serde_json::json!({"tls_profile":"chrome_136"})),
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
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 seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
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(move |_request: Request| async move {
|
||||
Json(json!({
|
||||
"request_id": "trace-gemini-chat-local-error-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"error": {
|
||||
"message": "quota reached",
|
||||
"status": "RESOURCE_EXHAUSTED"
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 27
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("client-gemini-chat-local-error-key")),
|
||||
sample_auth_snapshot("api-key-gemini-local-error-1", "user-gemini-local-error-1"),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/v1beta/models/gemini-2.5-pro:generateContent?key=client-gemini-chat-local-error-key&alt=sse"
|
||||
))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-gemini-chat-local-error-123")
|
||||
.body(
|
||||
"{\"contents\":[],\"generationConfig\":{\"temperature\":0.2},\"metadata\":{\"client\":\"desktop-gemini\"},\"toolConfig\":{\"functionCallingConfig\":{\"mode\":\"AUTO\"}}}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
|
||||
);
|
||||
let response_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
response_json,
|
||||
json!({
|
||||
"error": {
|
||||
"message": "quota reached",
|
||||
"status": "RESOURCE_EXHAUSTED"
|
||||
}
|
||||
})
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-gemini-chat-local-error-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
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();
|
||||
}
|
||||
|
||||
@@ -17,7 +17,7 @@ use super::{
|
||||
fn sample_video_task(
|
||||
id: &str,
|
||||
status: VideoTaskStatus,
|
||||
created_at_unix_secs: u64,
|
||||
created_at_unix_ms: u64,
|
||||
model: &str,
|
||||
user_id: &str,
|
||||
client_api_format: &str,
|
||||
@@ -54,20 +54,20 @@ fn sample_video_task(
|
||||
retry_count: 0,
|
||||
poll_interval_seconds: 10,
|
||||
next_poll_at_unix_secs: if status.is_active() {
|
||||
Some(created_at_unix_secs.saturating_add(10))
|
||||
Some(created_at_unix_ms.saturating_add(10))
|
||||
} else {
|
||||
None
|
||||
},
|
||||
poll_count: 0,
|
||||
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.saturating_add(30))
|
||||
Some(created_at_unix_ms.saturating_add(30))
|
||||
} else {
|
||||
None
|
||||
},
|
||||
updated_at_unix_secs: created_at_unix_secs.saturating_add(5),
|
||||
updated_at_unix_secs: created_at_unix_ms.saturating_add(5),
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
@@ -513,7 +513,7 @@ async fn gateway_cancels_openai_video_task_via_internal_async_task_endpoint() {
|
||||
"OpenAi": {
|
||||
"local_task_id": "task-openai-cancel",
|
||||
"upstream_task_id": "ext-video-task-123",
|
||||
"created_at_unix_secs": 100,
|
||||
"created_at_unix_ms": 100,
|
||||
"user_id": "user-1",
|
||||
"api_key_id": "api-key-task-openai-cancel",
|
||||
"model": "sora-2",
|
||||
@@ -675,7 +675,7 @@ async fn gateway_cancels_openai_video_task_via_internal_async_task_endpoint_with
|
||||
"OpenAi": {
|
||||
"local_task_id": "task-openai-cancel-direct",
|
||||
"upstream_task_id": "ext-video-task-123",
|
||||
"created_at_unix_secs": 100,
|
||||
"created_at_unix_ms": 100,
|
||||
"user_id": "user-1",
|
||||
"api_key_id": "api-key-task-openai-cancel",
|
||||
"model": "sora-2",
|
||||
|
||||
@@ -377,7 +377,7 @@ fn sample_request_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 {
|
||||
@@ -403,9 +403,9 @@ fn sample_request_candidate(
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
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")
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
]),
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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),
|
||||
},
|
||||
],
|
||||
));
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -201,6 +201,8 @@ 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(),
|
||||
@@ -223,9 +225,9 @@ 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")
|
||||
}
|
||||
|
||||
@@ -64,7 +64,7 @@ fn sample_gemini_video_task(
|
||||
next_poll_at_unix_secs: (!completed).then_some(124),
|
||||
poll_count: 0,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 123,
|
||||
created_at_unix_ms: 123,
|
||||
submitted_at_unix_secs: Some(123),
|
||||
completed_at_unix_secs: completed.then_some(124),
|
||||
updated_at_unix_secs: 124,
|
||||
|
||||
@@ -81,7 +81,7 @@ async fn gateway_reads_openai_video_task_via_data_read_side_without_hitting_publ
|
||||
next_poll_at_unix_secs: Some(124),
|
||||
poll_count: 1,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 123,
|
||||
created_at_unix_ms: 123,
|
||||
submitted_at_unix_secs: Some(123),
|
||||
completed_at_unix_secs: None,
|
||||
updated_at_unix_secs: 124,
|
||||
@@ -186,7 +186,7 @@ async fn gateway_reads_gemini_video_task_via_data_read_side_without_hitting_publ
|
||||
next_poll_at_unix_secs: None,
|
||||
poll_count: 3,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 223,
|
||||
created_at_unix_ms: 223,
|
||||
submitted_at_unix_secs: Some(223),
|
||||
completed_at_unix_secs: Some(224),
|
||||
updated_at_unix_secs: 224,
|
||||
|
||||
@@ -188,7 +188,7 @@ async fn gateway_executes_gemini_video_cancel_via_data_backed_local_follow_up_wi
|
||||
next_poll_at_unix_secs: Some(124),
|
||||
poll_count: 0,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 123,
|
||||
created_at_unix_ms: 123,
|
||||
submitted_at_unix_secs: Some(123),
|
||||
completed_at_unix_secs: None,
|
||||
updated_at_unix_secs: 123,
|
||||
@@ -551,7 +551,7 @@ async fn gateway_executes_gemini_video_cancel_via_reconstructed_data_backed_loca
|
||||
next_poll_at_unix_secs: Some(123),
|
||||
poll_count: 0,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 123,
|
||||
created_at_unix_ms: 123,
|
||||
submitted_at_unix_secs: Some(123),
|
||||
completed_at_unix_secs: None,
|
||||
updated_at_unix_secs: 123,
|
||||
|
||||
@@ -671,7 +671,7 @@ async fn gateway_executes_openai_video_remix_via_data_backed_local_follow_up_wit
|
||||
next_poll_at_unix_secs: None,
|
||||
poll_count: 1,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 123,
|
||||
created_at_unix_ms: 123,
|
||||
submitted_at_unix_secs: Some(123),
|
||||
completed_at_unix_secs: Some(124),
|
||||
updated_at_unix_secs: 124,
|
||||
@@ -683,7 +683,7 @@ async fn gateway_executes_openai_video_remix_via_data_backed_local_follow_up_wit
|
||||
"OpenAi": {
|
||||
"local_task_id": "task-local-123",
|
||||
"upstream_task_id": "ext-video-task-123",
|
||||
"created_at_unix_secs": 123,
|
||||
"created_at_unix_ms": 123,
|
||||
"user_id": "user-openai-video-remix-local-123",
|
||||
"api_key_id": "key-openai-video-remix-local-123",
|
||||
"model": "sora-2",
|
||||
|
||||
@@ -268,7 +268,7 @@ async fn gateway_executes_openai_video_delete_via_reconstructed_data_backed_loca
|
||||
next_poll_at_unix_secs: None,
|
||||
poll_count: 0,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 123,
|
||||
created_at_unix_ms: 123,
|
||||
submitted_at_unix_secs: Some(123),
|
||||
completed_at_unix_secs: Some(456),
|
||||
updated_at_unix_secs: 456,
|
||||
|
||||
@@ -45,7 +45,7 @@ fn sample_due_openai_task(upstream_base_url: &str) -> UpsertVideoTask {
|
||||
next_poll_at_unix_secs: Some(0),
|
||||
poll_count: 0,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 123,
|
||||
created_at_unix_ms: 123,
|
||||
submitted_at_unix_secs: Some(123),
|
||||
completed_at_unix_secs: None,
|
||||
updated_at_unix_secs: 123,
|
||||
@@ -57,7 +57,7 @@ fn sample_due_openai_task(upstream_base_url: &str) -> UpsertVideoTask {
|
||||
"OpenAi": {
|
||||
"local_task_id": "task-local-123",
|
||||
"upstream_task_id": "ext-video-task-123",
|
||||
"created_at_unix_secs": 123,
|
||||
"created_at_unix_ms": 123,
|
||||
"user_id": "user-video-poller-123",
|
||||
"api_key_id": "key-video-poller-123",
|
||||
"model": "sora-2",
|
||||
|
||||
@@ -287,7 +287,7 @@ async fn gateway_executes_openai_video_content_from_reconstructed_data_task_with
|
||||
next_poll_at_unix_secs: None,
|
||||
poll_count: 0,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs: 123,
|
||||
created_at_unix_ms: 123,
|
||||
submitted_at_unix_secs: Some(123),
|
||||
completed_at_unix_secs: Some(456),
|
||||
updated_at_unix_secs: 456,
|
||||
|
||||
Reference in New Issue
Block a user