mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 21:20:20 +08:00
Preserve terminal request candidate state
This commit is contained in:
@@ -1261,21 +1261,23 @@ async fn gateway_executes_openai_responses_sync_after_api_key_concurrency_wait_b
|
||||
*execution_runtime_hits.lock().expect("mutex should lock") >= 1
|
||||
})
|
||||
.await;
|
||||
let pending_deadline = tokio::time::Instant::now() + std::time::Duration::from_millis(5_000);
|
||||
let active_deadline = tokio::time::Instant::now() + std::time::Duration::from_millis(5_000);
|
||||
loop {
|
||||
let inflight_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-openai-cli-local-timeout-inflight-123")
|
||||
.await
|
||||
.expect("inflight request candidate trace should read");
|
||||
if inflight_candidates
|
||||
.iter()
|
||||
.any(|candidate| candidate.status == RequestCandidateStatus::Pending)
|
||||
{
|
||||
if inflight_candidates.iter().any(|candidate| {
|
||||
matches!(
|
||||
candidate.status,
|
||||
RequestCandidateStatus::Pending | RequestCandidateStatus::Streaming
|
||||
) && candidate.started_at_unix_ms.is_some()
|
||||
}) {
|
||||
break;
|
||||
}
|
||||
assert!(
|
||||
tokio::time::Instant::now() < pending_deadline,
|
||||
"inflight request candidate did not become pending"
|
||||
tokio::time::Instant::now() < active_deadline,
|
||||
"inflight request candidate did not become active: {inflight_candidates:?}"
|
||||
);
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
@@ -4,9 +4,9 @@ use std::sync::RwLock;
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::{
|
||||
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate,
|
||||
UpsertRequestCandidateRecord,
|
||||
request_candidate_lifecycle_would_regress, PublicHealthStatusCount, PublicHealthTimelineBucket,
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
|
||||
StoredRequestCandidate, UpsertRequestCandidateRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -311,6 +311,17 @@ impl RequestCandidateWriteRepository for InMemoryRequestCandidateRepository {
|
||||
})
|
||||
.cloned();
|
||||
|
||||
let preserve_existing_lifecycle = existing.as_ref().is_some_and(|row| {
|
||||
request_candidate_lifecycle_would_regress(row.status, candidate.status)
|
||||
});
|
||||
let merged_status = if preserve_existing_lifecycle {
|
||||
existing
|
||||
.as_ref()
|
||||
.map(|row| row.status)
|
||||
.unwrap_or(candidate.status)
|
||||
} else {
|
||||
candidate.status
|
||||
};
|
||||
let created_at_unix_ms = existing
|
||||
.as_ref()
|
||||
.map(|row| row.created_at_unix_ms)
|
||||
@@ -348,25 +359,41 @@ impl RequestCandidateWriteRepository for InMemoryRequestCandidateRepository {
|
||||
key_id: candidate
|
||||
.key_id
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.key_id.clone())),
|
||||
status: candidate.status,
|
||||
status: merged_status,
|
||||
skip_reason: candidate
|
||||
.skip_reason
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.skip_reason.clone())),
|
||||
is_cached: candidate
|
||||
.is_cached
|
||||
.unwrap_or_else(|| existing.as_ref().map(|row| row.is_cached).unwrap_or(false)),
|
||||
status_code: candidate
|
||||
.status_code
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.status_code)),
|
||||
error_type: candidate
|
||||
.error_type
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.error_type.clone())),
|
||||
error_message: candidate
|
||||
.error_message
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.error_message.clone())),
|
||||
latency_ms: candidate
|
||||
.latency_ms
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.latency_ms)),
|
||||
status_code: if preserve_existing_lifecycle {
|
||||
existing.as_ref().and_then(|row| row.status_code)
|
||||
} else {
|
||||
candidate
|
||||
.status_code
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.status_code))
|
||||
},
|
||||
error_type: if preserve_existing_lifecycle {
|
||||
existing.as_ref().and_then(|row| row.error_type.clone())
|
||||
} else {
|
||||
candidate
|
||||
.error_type
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.error_type.clone()))
|
||||
},
|
||||
error_message: if preserve_existing_lifecycle {
|
||||
existing.as_ref().and_then(|row| row.error_message.clone())
|
||||
} else {
|
||||
candidate
|
||||
.error_message
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.error_message.clone()))
|
||||
},
|
||||
latency_ms: if preserve_existing_lifecycle {
|
||||
existing.as_ref().and_then(|row| row.latency_ms)
|
||||
} else {
|
||||
candidate
|
||||
.latency_ms
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.latency_ms))
|
||||
},
|
||||
concurrent_requests: candidate
|
||||
.concurrent_requests
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.concurrent_requests)),
|
||||
@@ -383,9 +410,13 @@ impl RequestCandidateWriteRepository for InMemoryRequestCandidateRepository {
|
||||
started_at_unix_ms: candidate
|
||||
.started_at_unix_ms
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.started_at_unix_ms)),
|
||||
finished_at_unix_ms: candidate
|
||||
.finished_at_unix_ms
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.finished_at_unix_ms)),
|
||||
finished_at_unix_ms: if preserve_existing_lifecycle {
|
||||
existing.as_ref().and_then(|row| row.finished_at_unix_ms)
|
||||
} else {
|
||||
candidate
|
||||
.finished_at_unix_ms
|
||||
.or_else(|| existing.as_ref().and_then(|row| row.finished_at_unix_ms))
|
||||
},
|
||||
};
|
||||
|
||||
by_id.insert(stored.id.clone(), stored.clone());
|
||||
@@ -633,6 +664,84 @@ mod tests {
|
||||
assert_eq!(updated.started_at_unix_ms, Some(101));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_keeps_terminal_candidate_state_when_streaming_arrives_late() {
|
||||
let existing = StoredRequestCandidate::new(
|
||||
"cand-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
0,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
RequestCandidateStatus::Failed,
|
||||
None,
|
||||
false,
|
||||
Some(503),
|
||||
Some("upstream_error".to_string()),
|
||||
Some("retryable upstream failure".to_string()),
|
||||
Some(45),
|
||||
Some(1),
|
||||
Some(json!({"terminal": true})),
|
||||
None,
|
||||
100,
|
||||
Some(101),
|
||||
Some(145),
|
||||
)
|
||||
.expect("candidate should build");
|
||||
let repository = InMemoryRequestCandidateRepository::seed(vec![existing]);
|
||||
|
||||
let updated = repository
|
||||
.upsert(UpsertRequestCandidateRecord {
|
||||
id: "cand-1-late".to_string(),
|
||||
request_id: "req-1".to_string(),
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
candidate_index: 0,
|
||||
retry_index: 0,
|
||||
provider_id: None,
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
status: RequestCandidateStatus::Streaming,
|
||||
skip_reason: None,
|
||||
is_cached: None,
|
||||
status_code: Some(200),
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: Some(9_999),
|
||||
concurrent_requests: Some(2),
|
||||
extra_data: Some(json!({"late": true})),
|
||||
required_capabilities: None,
|
||||
created_at_unix_ms: None,
|
||||
started_at_unix_ms: Some(102),
|
||||
finished_at_unix_ms: None,
|
||||
})
|
||||
.await
|
||||
.expect("late update should succeed");
|
||||
|
||||
assert_eq!(updated.id, "cand-1");
|
||||
assert_eq!(updated.status, RequestCandidateStatus::Failed);
|
||||
assert_eq!(updated.status_code, Some(503));
|
||||
assert_eq!(updated.error_type.as_deref(), Some("upstream_error"));
|
||||
assert_eq!(
|
||||
updated.error_message.as_deref(),
|
||||
Some("retryable upstream failure")
|
||||
);
|
||||
assert_eq!(updated.latency_ms, Some(45));
|
||||
assert_eq!(updated.concurrent_requests, Some(2));
|
||||
assert_eq!(updated.finished_at_unix_ms, Some(145));
|
||||
assert_eq!(
|
||||
updated.extra_data,
|
||||
Some(json!({"terminal": true, "late": true}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delete_created_before_removes_oldest_matching_rows_up_to_limit() {
|
||||
let repository = InMemoryRequestCandidateRepository::seed(vec![
|
||||
|
||||
@@ -15,3 +15,23 @@ pub use memory::InMemoryRequestCandidateRepository;
|
||||
pub use mysql::MysqlRequestCandidateRepository;
|
||||
pub use postgres::SqlxRequestCandidateReadRepository;
|
||||
pub use sqlite::SqliteRequestCandidateRepository;
|
||||
|
||||
fn request_candidate_lifecycle_would_regress(
|
||||
existing: RequestCandidateStatus,
|
||||
incoming: RequestCandidateStatus,
|
||||
) -> bool {
|
||||
matches!(
|
||||
existing,
|
||||
RequestCandidateStatus::Success
|
||||
| RequestCandidateStatus::Failed
|
||||
| RequestCandidateStatus::Cancelled
|
||||
| RequestCandidateStatus::Skipped
|
||||
) && matches!(
|
||||
incoming,
|
||||
RequestCandidateStatus::Available
|
||||
| RequestCandidateStatus::Unused
|
||||
| RequestCandidateStatus::Pending
|
||||
| RequestCandidateStatus::Streaming
|
||||
) || existing == RequestCandidateStatus::Streaming
|
||||
&& incoming == RequestCandidateStatus::Pending
|
||||
}
|
||||
|
||||
@@ -4,9 +4,9 @@ use async_trait::async_trait;
|
||||
use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
|
||||
|
||||
use super::{
|
||||
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate,
|
||||
UpsertRequestCandidateRecord,
|
||||
request_candidate_lifecycle_would_regress, PublicHealthStatusCount, PublicHealthTimelineBucket,
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
|
||||
StoredRequestCandidate, UpsertRequestCandidateRecord,
|
||||
};
|
||||
use crate::driver::mysql::MysqlPool;
|
||||
use crate::error::SqlResultExt;
|
||||
@@ -525,26 +525,6 @@ fn merge_candidate(
|
||||
)
|
||||
}
|
||||
|
||||
fn request_candidate_lifecycle_would_regress(
|
||||
existing: RequestCandidateStatus,
|
||||
incoming: RequestCandidateStatus,
|
||||
) -> bool {
|
||||
matches!(
|
||||
existing,
|
||||
RequestCandidateStatus::Success
|
||||
| RequestCandidateStatus::Failed
|
||||
| RequestCandidateStatus::Cancelled
|
||||
| RequestCandidateStatus::Skipped
|
||||
) && matches!(
|
||||
incoming,
|
||||
RequestCandidateStatus::Available
|
||||
| RequestCandidateStatus::Unused
|
||||
| RequestCandidateStatus::Pending
|
||||
| RequestCandidateStatus::Streaming
|
||||
) || existing == RequestCandidateStatus::Streaming
|
||||
&& incoming == RequestCandidateStatus::Pending
|
||||
}
|
||||
|
||||
fn aggregate_timeline(
|
||||
candidates: Vec<StoredRequestCandidate>,
|
||||
since_unix_secs: u64,
|
||||
|
||||
@@ -4,9 +4,9 @@ use async_trait::async_trait;
|
||||
use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
|
||||
|
||||
use super::{
|
||||
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate,
|
||||
UpsertRequestCandidateRecord,
|
||||
request_candidate_lifecycle_would_regress, PublicHealthStatusCount, PublicHealthTimelineBucket,
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
|
||||
StoredRequestCandidate, UpsertRequestCandidateRecord,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::error::SqlResultExt;
|
||||
@@ -512,26 +512,6 @@ fn merge_candidate(
|
||||
)
|
||||
}
|
||||
|
||||
fn request_candidate_lifecycle_would_regress(
|
||||
existing: RequestCandidateStatus,
|
||||
incoming: RequestCandidateStatus,
|
||||
) -> bool {
|
||||
matches!(
|
||||
existing,
|
||||
RequestCandidateStatus::Success
|
||||
| RequestCandidateStatus::Failed
|
||||
| RequestCandidateStatus::Cancelled
|
||||
| RequestCandidateStatus::Skipped
|
||||
) && matches!(
|
||||
incoming,
|
||||
RequestCandidateStatus::Available
|
||||
| RequestCandidateStatus::Unused
|
||||
| RequestCandidateStatus::Pending
|
||||
| RequestCandidateStatus::Streaming
|
||||
) || existing == RequestCandidateStatus::Streaming
|
||||
&& incoming == RequestCandidateStatus::Pending
|
||||
}
|
||||
|
||||
fn aggregate_timeline(
|
||||
candidates: Vec<StoredRequestCandidate>,
|
||||
since_unix_secs: u64,
|
||||
|
||||
Reference in New Issue
Block a user