fix(store): persist retries after transport commit

Let Core decide replay eligibility after WebSocket preludes or HTTP
keepalives. Persist subsequent attempts while still enforcing the running
lifecycle and consecutive attempt numbers, preserving delivery facts and
the upstream send watermark.

Cover retry metadata, duplicate and skipped attempts, and terminal request
immutability with a PostgreSQL regression test.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-10-12 17:06:03 +08:00
co-authored by Codex
parent 9e02e52fec
commit 3ef4743e1a
3 changed files with 117 additions and 3 deletions
@@ -513,6 +513,8 @@ impl ModelRequestRepository for PgExecutionStore {
attempt: ModelRequestAttemptStart,
) -> StoreResult<u32> {
attempt.validate()?;
// 重试资格由 Core 的重放屏障裁决,已交付前导事件或心跳不代表禁止新 attempt
// 观测存储只校验运行态与连续序号,保留首次交付时间和客户端状态
// upstream_send_state 是请求级单调水位(sent > ambiguous > not_sent),
// 由 mark/finalize 抬升;开启新 attempt 不得把已持久化的 sent 重置回
// not_sent——崩溃恢复的终态写回会原样继承本列
@@ -539,7 +541,7 @@ impl ModelRequestRepository for PgExecutionStore {
'capacityUsedSlots', coalesce($10::bigint, observed.used_slots),
'capacityTotalSlots', coalesce($11::bigint, observed.total_slots))))
from observed
where model_requests.id = observed.id and outcome = 'running' and downstream_committed_at is null
where model_requests.id = observed.id and outcome = 'running'
and $8 = attempt_count + 1
returning attempt_count",
)
@@ -573,7 +575,7 @@ impl ModelRequestRepository for PgExecutionStore {
source: None,
entity: ENTITY,
id: attempt.model_request_id,
kind: ConflictKind::DownstreamAlreadyCommitted,
kind: ConflictKind::InvalidTransition,
})?;
u32::try_from(count)
.map_err(|source| invalid("attempt_count is invalid").with_source(source))
@@ -15,7 +15,6 @@ pub enum ConflictKind {
StaleRevision,
DuplicateName,
AlreadyFinalized,
DownstreamAlreadyCommitted,
RequestNotRunning,
InvalidTransition,
LeaseLost,
@@ -2241,6 +2241,119 @@ async fn shared_core_usage_and_timings_preserve_fields_and_reject_invalid_phases
database.close().await;
}
#[tokio::test]
async fn retry_after_transport_commit_preserves_delivery_facts_and_attempt_order() {
let Some(database) = TestDatabase::create("retry_after_commit").await else {
return;
};
let store = PgExecutionStore::new(database.pool.clone());
for (transport, status) in [("websocket", None), ("http_sse", Some(200))] {
let id = format!("req_retry_after_{transport}_commit");
let mut request = accepted_request(&id);
request.client_transport = transport.to_owned();
store.create_model_request(request).await.unwrap();
let mut attempt = ModelRequestAttemptStart {
model_request_id: id.clone(),
attempt_count: 1,
provider_kind: "openai".to_owned(),
account_snapshot: None,
provider_account_ref: Some("account_before_retry".to_owned()),
upstream_model_id: Some("model".to_owned()),
upstream_transport: "websocket".to_owned(),
http_version: Some("HTTP/1.1".to_owned()),
account_selection_wait_ms: Some(5),
capacity_used_slots: Some(1),
capacity_total_slots: Some(10),
};
store
.begin_model_request_attempt(attempt.clone())
.await
.unwrap();
store
.mark_upstream_send_state(&id, gateway_store::postgres::UpstreamSendState::Sent)
.await
.unwrap();
// 可重放的 WS 前导事件或 HTTP 心跳先交付,不能替代 Core 的业务重放屏障
let committed_at = Utc::now();
ModelRequestRepository::mark_downstream_committed(&store, &id, committed_at, status)
.await
.unwrap();
let committed = stored_row(&database.pool, &id).await;
attempt.attempt_count = 2;
attempt.provider_account_ref = Some("account_after_retry".to_owned());
attempt.upstream_transport = "http_sse".to_owned();
attempt.http_version = Some("HTTP/2".to_owned());
attempt.account_selection_wait_ms = Some(7);
attempt.capacity_used_slots = Some(2);
assert_eq!(
store
.begin_model_request_attempt(attempt.clone())
.await
.unwrap(),
2
);
let retried = stored_row(&database.pool, &id).await;
assert_eq!(
retried["downstream_committed_at"],
committed["downstream_committed_at"]
);
assert_eq!(
retried["client_status_code"],
committed["client_status_code"]
);
assert_eq!(retried["upstream_send_state"], "sent");
assert_eq!(retried["provider_account_ref"], "account_after_retry");
assert_eq!(retried["upstream_transport"], "http_sse");
assert_eq!(
retried["request_observation_json"]["scheduling"]["accountSelectionWaitMs"],
12
);
assert_eq!(
retried["request_observation_json"]["scheduling"]["capacityUsedSlots"],
2
);
assert_eq!(
retried["request_observation_json"]["transport"]["httpVersion"],
"HTTP/2"
);
// 已接收的序号与跳号都不能覆盖新 attempt 的账号和计时
for count in [2, 4] {
attempt.attempt_count = count;
let error = store
.begin_model_request_attempt(attempt.clone())
.await
.unwrap_err();
assert!(matches!(
error,
gateway_store::StoreError::Conflict {
kind: gateway_store::ConflictKind::InvalidTransition,
..
}
));
assert_eq!(stored_row(&database.pool, &id).await, retried);
}
let mut finalization = successful_core_finalization(&id);
finalization.attempt_count = 2;
finalization.downstream_committed_at = Some(committed_at.into());
finalization.client_status_code = status;
finalization.upstream_transport = Some("http_sse".to_owned());
ExecutionStore::finalize_model_request(&store, finalization)
.await
.unwrap();
let terminal = stored_row(&database.pool, &id).await;
assert_eq!(terminal["outcome"], "succeeded");
assert_eq!(terminal["attempt_count"], 2);
assert_eq!(
terminal["request_observation_json"]["scheduling"],
retried["request_observation_json"]["scheduling"]
);
attempt.attempt_count = 3;
assert!(store.begin_model_request_attempt(attempt).await.is_err());
assert_eq!(stored_row(&database.pool, &id).await, terminal);
}
database.close().await;
}
#[tokio::test]
async fn concurrent_retry_updates_keep_request_context_and_clear_previous_transport() {
let Some(database) = TestDatabase::create("observation_retry_cas").await else {