fix(data): repair request candidate epoch created_at (#343)

Co-authored-by: fawney19 <elky0401@gmail.com>
This commit is contained in:
RWDai
2026-04-30 09:23:21 +08:00
committed by GitHub
parent fa47e8a3c0
commit 2b5247b9a8
5 changed files with 255 additions and 161 deletions

View File

@@ -300,19 +300,17 @@ mod tests {
}
#[test]
fn admin_user_api_formats_reject_retired_signatures() {
for retired in [
"anthropic:messages",
"claude:chat",
"claude:cli",
"openai:cli",
"openai:compact",
"gemini:chat",
"gemini:cli",
fn admin_user_api_formats_reject_unsupported_signatures() {
for unsupported in [
"claude",
"openai",
"unknown:chat",
"openai:unknown",
"gemini:generate",
] {
assert!(
normalize_admin_user_api_formats(Some(vec![retired.to_string()])).is_err(),
"{retired} should be rejected"
normalize_admin_user_api_formats(Some(vec![unsupported.to_string()])).is_err(),
"{unsupported} should be rejected"
);
}
}

View File

@@ -25,6 +25,27 @@ fn deep_nested_metadata(levels: usize) -> serde_json::Value {
current
}
fn run_async_test_on_large_stack<F>(name: &'static str, future: F)
where
F: std::future::Future<Output = ()> + Send + 'static,
{
let handle = std::thread::Builder::new()
.name(name.to_string())
.stack_size(8 * 1024 * 1024)
.spawn(move || {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("tokio runtime should build")
.block_on(future);
})
.expect("large-stack test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
async fn wait_for_usage_status<T>(
repository: &T,
request_id: &str,
@@ -232,8 +253,15 @@ async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage() {
#[test]
fn gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage() {
run_async_test_on_large_stack(
"gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage",
gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage_impl(),
);
}
async fn gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage_impl() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
@@ -812,8 +840,15 @@ async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exha
assert_eq!(stored_candidates[0].status_code, Some(503));
}
#[tokio::test]
async fn gateway_records_failed_usage_for_claude_runtime_miss_without_execution_exhaustion() {
#[test]
fn gateway_records_failed_usage_for_claude_runtime_miss_without_execution_exhaustion() {
run_async_test_on_large_stack(
"gateway_records_failed_usage_for_claude_runtime_miss_without_execution_exhaustion",
gateway_records_failed_usage_for_claude_runtime_miss_without_execution_exhaustion_impl(),
);
}
async fn gateway_records_failed_usage_for_claude_runtime_miss_without_execution_exhaustion_impl() {
fn sample_claude_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
@@ -1357,8 +1392,15 @@ async fn gateway_preserves_stream_usage_when_max_response_body_size_truncates_ca
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_skipped() {
#[test]
fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_skipped() {
run_async_test_on_large_stack(
"gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_skipped",
gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_skipped_impl(),
);
}
async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_skipped_impl() {
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
@@ -1670,9 +1712,9 @@ async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_s
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_keeps_failed_usage_request_capture_lightweight_for_large_local_claude_cli_runtime_miss(
) {
#[test]
fn gateway_keeps_failed_usage_request_capture_lightweight_for_large_local_claude_cli_runtime_miss()
{
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
@@ -1810,122 +1852,124 @@ async fn gateway_keeps_failed_usage_request_capture_lightweight_for_large_local_
.expect("key transport should build")
}
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
run_async_test_on_large_stack("large-local-claude-cli-runtime-miss", async move {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let execution_runtime = Router::new();
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-claude-cli-usage-local-miss-large")),
sample_auth_snapshot(
"api-key-claude-cli-usage-local-miss-large-1",
"user-claude-cli-usage-local-miss-large-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 execution_runtime = Router::new();
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-claude-cli-usage-local-miss-large")),
sample_auth_snapshot(
"api-key-claude-cli-usage-local-miss-large-1",
"user-claude-cli-usage-local-miss-large-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 (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state =
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
let (execution_runtime_url, execution_runtime_handle) =
start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let request_body = serde_json::to_string(&json!({
"model": "gpt-5.4",
"messages": [{
"role": "user",
"content": "x".repeat(128 * 1024)
}],
"metadata": deep_nested_metadata(96)
}))
.expect("request should encode");
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/messages?beta=true"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-claude-cli-usage-local-miss-large",
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
.header(
TRACE_ID_HEADER,
"trace-claude-cli-usage-local-miss-large-123",
)
.body(request_body)
.send()
.await
.expect("request should complete");
let request_body = serde_json::to_string(&json!({
"model": "gpt-5.4",
"messages": [{
"role": "user",
"content": "x".repeat(128 * 1024)
}],
"metadata": deep_nested_metadata(96)
}))
.expect("request should encode");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
response
.headers()
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok()),
Some("all_candidates_skipped")
);
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/messages?beta=true"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-claude-cli-usage-local-miss-large",
)
.header(
TRACE_ID_HEADER,
let stored_usage = wait_for_usage_status(
usage_repository.as_ref(),
"trace-claude-cli-usage-local-miss-large-123",
"failed",
)
.body(request_body)
.send()
.await
.expect("request should complete");
.await;
assert_eq!(stored_usage.status, "failed");
assert_eq!(
stored_usage.request_body_state,
Some(UsageBodyCaptureState::Inline)
);
assert_eq!(
stored_usage
.request_body
.as_ref()
.and_then(|value| value.get("model"))
.and_then(|value| value.as_str()),
Some("gpt-5.4")
);
assert!(stored_usage.provider_request_body.is_none());
assert_eq!(
stored_usage
.request_metadata
.as_ref()
.and_then(|value| value.get("trace_id"))
.and_then(|value| value.as_str()),
Some("trace-claude-cli-usage-local-miss-large-123")
);
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
response
.headers()
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok()),
Some("all_candidates_skipped")
);
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-claude-cli-usage-local-miss-large-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Skipped);
assert_eq!(
stored_candidates[0].skip_reason.as_deref(),
Some("format_conversion_disabled")
);
let stored_usage = wait_for_usage_status(
usage_repository.as_ref(),
"trace-claude-cli-usage-local-miss-large-123",
"failed",
)
.await;
assert_eq!(stored_usage.status, "failed");
assert_eq!(
stored_usage.request_body_state,
Some(UsageBodyCaptureState::Inline)
);
assert_eq!(
stored_usage
.request_body
.as_ref()
.and_then(|value| value.get("model"))
.and_then(|value| value.as_str()),
Some("gpt-5.4")
);
assert!(stored_usage.provider_request_body.is_none());
assert_eq!(
stored_usage
.request_metadata
.as_ref()
.and_then(|value| value.get("trace_id"))
.and_then(|value| value.as_str()),
Some("trace-claude-cli-usage-local-miss-large-123")
);
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-claude-cli-usage-local-miss-large-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Skipped);
assert_eq!(
stored_candidates[0].skip_reason.as_deref(),
Some("format_conversion_disabled")
);
gateway_handle.abort();
execution_runtime_handle.abort();
gateway_handle.abort();
execution_runtime_handle.abort();
});
}

View File

@@ -219,7 +219,14 @@ VALUES (
$19,
$20,
$21,
TO_TIMESTAMP(COALESCE($22, 0) / 1000.0),
COALESCE(
CASE
WHEN $22 IS NOT NULL AND $22 > 1000.0 THEN TO_TIMESTAMP($22 / 1000.0)
END,
TO_TIMESTAMP($23 / 1000.0),
TO_TIMESTAMP($24 / 1000.0),
NOW()
),
TO_TIMESTAMP($23 / 1000.0),
TO_TIMESTAMP($24 / 1000.0)
)
@@ -249,6 +256,11 @@ DO UPDATE SET
ELSE EXCLUDED.extra_data
END,
required_capabilities = COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities),
created_at = CASE
WHEN request_candidates.created_at <= TO_TIMESTAMP(1)
THEN EXCLUDED.created_at
ELSE request_candidates.created_at
END,
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
RETURNING
@@ -740,9 +752,21 @@ fn to_i32_u64(value: u64) -> Result<i32, DataLayerError> {
#[cfg(test)]
mod tests {
use super::SqlxRequestCandidateReadRepository;
use super::{SqlxRequestCandidateReadRepository, UPSERT_SQL};
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[test]
fn upsert_sql_does_not_default_missing_or_epoch_created_at_to_epoch() {
assert!(!UPSERT_SQL.contains("COALESCE($22, 0)"));
assert!(UPSERT_SQL.contains("WHEN $22 IS NOT NULL AND $22 > 1000.0"));
assert!(UPSERT_SQL.contains("TO_TIMESTAMP($22 / 1000.0)"));
assert!(UPSERT_SQL.contains("TO_TIMESTAMP($23 / 1000.0)"));
assert!(UPSERT_SQL.contains("TO_TIMESTAMP($24 / 1000.0)"));
assert!(UPSERT_SQL.contains("NOW()"));
assert!(UPSERT_SQL.contains("request_candidates.created_at <= TO_TIMESTAMP(1)"));
assert!(UPSERT_SQL.contains("THEN EXCLUDED.created_at"));
}
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {

View File

@@ -156,7 +156,7 @@ mod tests {
"claude",
));
assert!(!provider_matches_allowed_value(
"anthropic",
"vendor-x",
"provider-1",
"Other",
"claude",
@@ -164,11 +164,11 @@ mod tests {
assert!(!provider_matches_allowed_value(
"claude",
"provider-1",
"Anthropic",
"OtherVendor",
"custom",
));
assert!(!provider_matches_allowed_value(
"anthropic:messages",
"provider-1:extra",
"provider-1",
"Other",
"claude",
@@ -201,38 +201,22 @@ mod tests {
}
#[test]
fn api_format_allowed_value_rejects_retired_aliases() {
assert!(!api_format_matches_allowed_value(
"anthropic:messages",
"claude:messages"
));
assert!(!api_format_matches_allowed_value(
"claude:chat",
"claude:messages"
));
assert!(!api_format_matches_allowed_value(
"claude:cli",
"claude:messages"
));
assert!(!api_format_matches_allowed_value(
"openai:cli",
"openai:responses"
));
assert!(!api_format_matches_allowed_value(
"openai:compact",
"openai:responses:compact"
));
assert!(!api_format_matches_allowed_value(
"gemini:chat",
"gemini:generate_content"
));
fn api_format_allowed_value_matches_current_signatures_only() {
assert!(api_format_matches_allowed_value(
"CLAUDE:MESSAGES",
"claude:messages"
));
assert!(api_format_matches_allowed_value(
"openai:responses",
"openai:responses"
));
assert!(!api_format_matches_allowed_value(
"openai:responses",
"claude:messages"
));
assert!(!api_format_matches_allowed_value(
"claude:messages:extra",
"claude:messages"
));
}
}

View File

@@ -438,6 +438,10 @@ pub fn build_report_request_candidate_status_record(
let finished_at_unix_ms = finished_at_unix_ms
.or(slot.finished_at_unix_ms)
.or_else(|| is_terminal_candidate_status(status).then_some(terminal_unix_secs));
let created_at_unix_ms = non_epoch_unix_ms(slot.created_at_unix_ms)
.or_else(|| started_at_unix_ms.and_then(non_epoch_unix_ms))
.or_else(|| finished_at_unix_ms.and_then(non_epoch_unix_ms))
.unwrap_or(terminal_unix_secs);
UpsertRequestCandidateRecord {
id: slot.id,
@@ -461,12 +465,16 @@ pub fn build_report_request_candidate_status_record(
concurrent_requests: None,
extra_data: slot.extra_data,
required_capabilities: None,
created_at_unix_ms: Some(slot.created_at_unix_ms),
created_at_unix_ms: Some(created_at_unix_ms),
started_at_unix_ms,
finished_at_unix_ms,
}
}
fn non_epoch_unix_ms(value: u64) -> Option<u64> {
(value > 1000).then_some(value)
}
pub fn finalize_execution_request_candidate_report_context(
report_context: Value,
candidate_id: &str,
@@ -1072,7 +1080,43 @@ mod tests {
assert_eq!(record.started_at_unix_ms, Some(123));
assert_eq!(record.finished_at_unix_ms, Some(123));
assert_eq!(record.created_at_unix_ms, Some(10));
assert_eq!(record.created_at_unix_ms, Some(123));
assert_eq!(record.status, RequestCandidateStatus::Success);
}
#[test]
fn report_request_candidate_status_record_repairs_epoch_created_at() {
let record =
build_report_request_candidate_status_record(ReportRequestCandidateStatusRecordInput {
slot: SchedulerResolvedReportRequestCandidateSlot {
id: "cand-epoch".to_string(),
request_id: "req-epoch".to_string(),
user_id: None,
api_key_id: None,
candidate_index: 0,
retry_index: 0,
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
extra_data: None,
created_at_unix_ms: 0,
started_at_unix_ms: None,
finished_at_unix_ms: None,
},
status_update: SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
status_code: Some(200),
error_type: None,
error_message: None,
latency_ms: Some(12),
started_at_unix_ms: Some(2_000),
finished_at_unix_ms: Some(3_000),
},
now_unix_ms: 4_000,
});
assert_eq!(record.created_at_unix_ms, Some(2_000));
assert_eq!(record.started_at_unix_ms, Some(2_000));
assert_eq!(record.finished_at_unix_ms, Some(3_000));
}
}