diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index 449993b98..60bdb4efb 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -210,6 +210,8 @@ struct CurrentAttempt { index: NonZeroU32, started_at: SystemTime, send_observed: bool, + /// Provider 已为当前 attempt 交付可供同请求恢复的不透明状态检查点。 + session_checkpoint_observed: bool, response_observation: Option, } @@ -570,6 +572,9 @@ where return Err(EngineError::Deadline); } PollBoundary::Item(Some(Ok(mut event))) => { + if let Some(state) = event.session_update() { + self.observe_session_checkpoint(state); + } if let Some(observation) = event.take_observation() { self.observe_response(observation); } @@ -931,6 +936,7 @@ where index: next_attempt, started_at: SystemTime::now(), send_observed: false, + session_checkpoint_observed: false, response_observation: None, }); Ok(None) @@ -972,6 +978,17 @@ where self.mark_send_observed().await; } + fn observe_session_checkpoint(&mut self, state: &ProviderSessionState) { + let Some(current) = self.current.as_mut() else { + return; + }; + if state.provider() != current.metadata.provider().as_str() { + return; + } + current.session_checkpoint_observed = true; + self.operation.set_provider_session_state(state.clone()); + } + async fn mark_send_observed(&mut self) { let Some(current) = self.current.as_mut() else { return; @@ -1071,6 +1088,11 @@ where .await; } let provider_proved_replay_safe = provider_proved_replay_safe(&error); + // Ambiguous 仍是默认 fail-closed。唯一例外是 Provider 在同一 attempt 中先 + // 交付了不透明恢复检查点,并明确证明从该检查点重放安全;该例外只用于 + // 固定原账号的 transport recovery,不放开普通换号或 continuation 重放。 + let checkpointed_transport_replay = + current.session_checkpoint_observed && error.replay_is_safe(); let continuation_retry = self.prepare_continuation_retry( ¤t, &error, @@ -1093,14 +1115,16 @@ where delay, }) if self.downstream_committed_at.is_none() && !self.delivery_pending - && attempt_send_state != UpstreamSendState::Ambiguous => + && (attempt_send_state != UpstreamSendState::Ambiguous + || checkpointed_transport_replay) => { Some((AttemptTransport::Retry(retry_index), delay)) } Some(crate::error::PreDeliveryRetry::SameAccountTransportFallback) if self.downstream_committed_at.is_none() && !self.delivery_pending - && attempt_send_state != UpstreamSendState::Ambiguous => + && (attempt_send_state != UpstreamSendState::Ambiguous + || checkpointed_transport_replay) => { Some((AttemptTransport::Fallback, Duration::ZERO)) } diff --git a/backend/crates/gateway-core/src/error.rs b/backend/crates/gateway-core/src/error.rs index d245030e4..ed810e1e5 100644 --- a/backend/crates/gateway-core/src/error.rs +++ b/backend/crates/gateway-core/src/error.rs @@ -623,13 +623,21 @@ impl ProviderError { .get_or_insert_with(|| Box::new(ProviderErrorFailureObservation::default())) } - /// 标记 Provider 已证明本次拒绝没有执行生成,可在下游提交前重放。 + /// 标记 Provider 已证明本次失败可在下游提交前安全重放。 + /// + /// 证明可以来自明确的“未执行”拒绝,也可以来自同 attempt 已交付、可恢复 + /// 同一业务请求的状态检查点;后者仍须由 Core 的同账号恢复门禁单独校验。 #[must_use] pub const fn with_replay_safe(mut self) -> Self { self.replay_safe = true; self } + /// 原地标记 Provider 已证明本次失败可从已交付的状态检查点安全恢复。 + pub const fn set_replay_safe(&mut self) { + self.replay_safe = true; + } + /// 允许 Core 仅在客户端尚未收到任何事件时执行一次受预算约束的换号恢复。 /// /// 该标记不证明上游未执行,也不等价于 [`Self::with_replay_safe`];它只表达 diff --git a/backend/crates/gateway-core/src/event.rs b/backend/crates/gateway-core/src/event.rs index 768db7898..cc615d5b7 100644 --- a/backend/crates/gateway-core/src/event.rs +++ b/backend/crates/gateway-core/src/event.rs @@ -763,7 +763,8 @@ impl ProviderEvent { } } - /// 把本事件标记为 Provider 连接内状态的提交边界。 + /// 附加 Provider 私有状态检查点;Core 可将其用于同请求恢复,协议连接可在 + /// terminal 事件上持有它供下一轮使用。 pub fn attach_session_update(&mut self, state: ProviderSessionState) { self.session_update = Some(Box::new(state)); } diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index 7216c0165..79bcb379a 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -270,6 +270,7 @@ enum Script { struct ScriptedProvider { scripts: Mutex>, contexts: Mutex>, + operations: Mutex>, } impl ScriptedProvider { @@ -277,6 +278,7 @@ impl ScriptedProvider { Self { scripts: Mutex::new(scripts.into()), contexts: Mutex::new(Vec::new()), + operations: Mutex::new(Vec::new()), } } } @@ -302,6 +304,10 @@ impl Provider for ScriptedProvider { request: ProviderRequest, context: AttemptContext, ) -> Result { + self.operations + .lock() + .expect("operations lock") + .push(request.operation().clone()); self.contexts.lock().expect("contexts lock").push(context); let script = self .scripts @@ -2838,14 +2844,25 @@ fn ambiguous_send_state_stops_retry() { fn ambiguous_pre_delivery_retry_marker_does_not_rotate_account() { let operation = generate_operation(); let route_plan = plan(&operation); + let checkpoint = ProviderSessionState::new( + "openai", + Map::from_iter([("checkpoint".to_owned(), json!("opaque-state"))]), + ) + .expect("provider checkpoint"); + let mut checkpoint_event = ProviderEvent::observation(ProviderResponseObservation::new( + UpstreamTransport::new("websocket").expect("transport"), + )); + checkpoint_event.attach_session_update(checkpoint); let (coordinator, store, provider) = coordinator(vec![ - Script::Stream { + Script::ObservedStream { account_id: "acct_first", - items: vec![Err(ProviderError::new( - ProviderErrorKind::Transport, - UpstreamSendState::Ambiguous, - ) - .with_pre_delivery_retry())], + items: vec![ + Ok(checkpoint_event), + Err( + ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous) + .with_pre_delivery_retry(), + ), + ], }, Script::Stream { account_id: "acct_second", @@ -2874,6 +2891,70 @@ fn ambiguous_pre_delivery_retry_marker_does_not_rotate_account() { assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed); } +#[test] +fn ambiguous_transport_retry_uses_provider_session_checkpoint_on_same_account() { + let operation = generate_operation(); + let route_plan = plan(&operation); + let checkpoint = ProviderSessionState::new( + "openai", + Map::from_iter([("checkpoint".to_owned(), json!("opaque-state"))]), + ) + .expect("provider checkpoint"); + let mut checkpoint_event = ProviderEvent::observation(ProviderResponseObservation::new( + UpstreamTransport::new("websocket").expect("transport"), + )); + checkpoint_event.attach_session_update(checkpoint.clone()); + let retry_index = NonZeroU32::new(1).expect("retry index"); + let (coordinator, store, provider) = coordinator(vec![ + Script::ObservedStream { + account_id: "acct_first", + items: vec![ + Ok(checkpoint_event), + Err( + ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous) + .with_replay_safe() + .with_pre_delivery_transport_retry(retry_index, Duration::ZERO), + ), + ], + }, + Script::Stream { + account_id: "acct_first", + items: complete_stream(None), + }, + ]); + + let mut session = block_on(coordinator.start( + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, + None, + None, + CancellationToken::new(), + )) + .expect("start execution"); + block_on(session.collect_uncommitted()).expect("checkpointed retry succeeds"); + block_on(session.commit_downstream(Some(200))).expect("commit winning response"); + + let original = ProviderAccountId::new("acct_first").expect("account id"); + let contexts = provider.contexts.lock().expect("contexts lock"); + assert_eq!(contexts.len(), 2); + assert_eq!( + contexts[1].transport(), + AttemptTransport::Retry(retry_index) + ); + assert_eq!(contexts[1].required_account(), Some(&original)); + drop(contexts); + let operations = provider.operations.lock().expect("operations lock"); + assert_eq!( + operations[1].provider_session_state("openai"), + Some(&checkpoint) + ); + let state = store.state.lock().expect("store lock"); + assert_eq!(state.intermediate_failures, 1); + assert_eq!(state.finalizations[0].attempt_count, 2); + assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded); +} + #[test] fn pre_delivery_transport_fallback_retries_the_same_account_with_fallback_transport() { let operation = generate_operation(); diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index 91461dd6a..c22c30ad4 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -186,6 +186,13 @@ pub(super) struct OpenAiSessionState { #[serde(default)] pub(super) client_turn_id: Option, pub(super) continuation_scope: OpenAiContinuationScope, + /// 该状态来自同一业务请求的 metadata 检查点,只允许被紧随其后的恢复 attempt 消费。 + #[serde(default, skip_serializing_if = "is_false")] + pub(super) retry_checkpoint: bool, +} + +const fn is_false(value: &bool) -> bool { + !*value } #[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -232,6 +239,26 @@ pub(super) fn encode_openai_session_state( .map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::Sent)) } +fn encode_openai_session_capture( + capture: &OpenAiSessionCapture, + retry_checkpoint: bool, +) -> Result { + let Some(continuation_scope) = capture.continuation_scope else { + return Err(provider_error( + ProviderErrorKind::Protocol, + UpstreamSendState::Sent, + )); + }; + encode_openai_session_state(OpenAiSessionState { + account_id: capture.account_id.clone(), + conversation_id: capture.conversation_id.clone(), + turn_state: capture.turn_state.clone(), + client_turn_id: capture.client_turn_id.clone(), + continuation_scope, + retry_checkpoint, + }) +} + pub(super) fn attach_openai_session_update( events: &mut [ProviderEvent], capture: &mut Option, @@ -245,16 +272,7 @@ pub(super) fn attach_openai_session_update( let Some(capture) = capture.take() else { return; }; - let Some(continuation_scope) = capture.continuation_scope else { - return; - }; - let Ok(update) = encode_openai_session_state(OpenAiSessionState { - account_id: capture.account_id, - conversation_id: capture.conversation_id, - turn_state: capture.turn_state, - client_turn_id: capture.client_turn_id, - continuation_scope, - }) else { + let Ok(update) = encode_openai_session_capture(&capture, false) else { return; }; events[terminal_index].attach_session_update(update); @@ -563,6 +581,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { WebSocketRecoveryContext { policy, requirement: request_transport_requirement, + replay_checkpoint_available: false, retry_count: websocket_retry_count, max_retries: stream_max_retries, request_id: context.request_id().as_str(), @@ -594,6 +613,9 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { Err(failure.error)?; return; } + let mut turn_state_checkpoint_observed = response.transport + == CodexBackendTransport::WebSocket + && response.turn_state.is_some(); if let Some(capture) = session_capture.as_mut() { capture.continuation_scope = Some(if capture.response_store { OpenAiContinuationScope::Persisted @@ -707,6 +729,49 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { continue; } Err(mut failure) => { + let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await; + let rate_limits_changed = if updates.is_empty() { + false + } else { + passive_quota_observation.observe(&updates); + let update_headers = rate_limit_update_headers(&updates); + observation_state.merge_rate_limit_headers(&update_headers) + }; + let turn_state_merge = merge_turn_state_update( + turn_state_updates.as_ref(), + &mut session_capture, + &mut observation_state, + ) + .await; + turn_state_checkpoint_observed |= turn_state_merge.is_some(); + let retry_checkpoint_update = if turn_state_checkpoint_observed + && failure.websocket_checkpoint_retryable + { + session_capture + .as_ref() + .and_then(|capture| encode_openai_session_capture(capture, true).ok()) + } else { + None + }; + let mut observation_event = if rate_limits_changed + || turn_state_merge.is_some() + || retry_checkpoint_update.is_some() + { + observation_state + .observation() + .map(ProviderEvent::observation) + } else { + None + }; + if let (Some(event), Some(update)) = + (observation_event.as_mut(), retry_checkpoint_update) + { + event.attach_session_update(update); + } + let replay_checkpoint_available = observation_event + .as_ref() + .and_then(ProviderEvent::session_update) + .is_some(); if failure.websocket_transport_retryable && response_transport == CodexBackendTransport::WebSocket && (matches!( @@ -719,6 +784,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { WebSocketRecoveryContext { policy: WebSocketFailurePolicy::Budgeted, requirement: request_transport_requirement, + replay_checkpoint_available, retry_count: websocket_retry_count, max_retries: stream_max_retries, request_id: context.request_id().as_str(), @@ -731,24 +797,8 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { }, ); } - let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await; - let rate_limits_changed = if updates.is_empty() { - false - } else { - passive_quota_observation.observe(&updates); - let update_headers = rate_limit_update_headers(&updates); - observation_state.merge_rate_limit_headers(&update_headers) - }; - let turn_state_changed = merge_turn_state_update( - turn_state_updates.as_ref(), - &mut session_capture, - &mut observation_state, - ) - .await; - if (rate_limits_changed || turn_state_changed) - && let Some(observation) = observation_state.observation() - { - yield ProviderEvent::observation(observation); + if let Some(event) = observation_event { + yield event; } if allows_account_state_mutation { synchronize_passive_quota( @@ -772,12 +822,14 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { passive_quota_observation.observe(&updates); observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates)) }; - let turn_state_changed = merge_turn_state_update( + let turn_state_merge = merge_turn_state_update( turn_state_updates.as_ref(), &mut session_capture, &mut observation_state, ) .await; + turn_state_checkpoint_observed |= turn_state_merge.is_some(); + let turn_state_changed = turn_state_merge.unwrap_or(false); let first_event_changed = observation_state.observe_stream_chunk(&chunk, output_started_at); let chunk_len = chunk.len(); @@ -959,7 +1011,8 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { &mut session_capture, &mut observation_state, ) - .await; + .await + .unwrap_or(false); attach_openai_session_update(&mut events, &mut session_capture); let completed = events .iter() @@ -1030,15 +1083,11 @@ async fn merge_turn_state_update( updates: Option<&CodexTurnStateUpdate>, session_capture: &mut Option, observation_state: &mut OpenAiResponseObservationState, -) -> bool { - let Some(updates) = updates else { - return false; - }; - let Some(turn_state) = updates.lock().await.take() else { - return false; - }; +) -> Option { + let updates = updates?; + let turn_state = updates.lock().await.take()?; if let Some(capture) = session_capture.as_mut() { capture.turn_state = Some(turn_state.clone()); } - observation_state.merge_client_header("x-codex-turn-state", &turn_state) + Some(observation_state.merge_client_header("x-codex-turn-state", &turn_state)) } diff --git a/backend/crates/providers/openai/src/provider/failure.rs b/backend/crates/providers/openai/src/provider/failure.rs index 22c7eb35a..9289cc9a3 100644 --- a/backend/crates/providers/openai/src/provider/failure.rs +++ b/backend/crates/providers/openai/src/provider/failure.rs @@ -39,6 +39,8 @@ pub fn openai_failure_affects_account_score(error: &ProviderError) -> bool { pub(super) struct MappedProviderFailure { pub(super) error: ProviderError, pub(super) websocket_transport_retryable: bool, + /// 兼容少数链路在 metadata 后、终态前关闭;仅在捕获 turn-state 时允许安全恢复。 + pub(super) websocket_checkpoint_retryable: bool, pub(super) account_failure: Option, /// 原始上游错误描述,仅在凭据错误状态下持久化。 pub(super) error_message: Option, @@ -54,6 +56,7 @@ impl MappedProviderFailure { Self { error, websocket_transport_retryable: false, + websocket_checkpoint_retryable: false, account_failure: None, error_message: None, cyber_policy_failure: false, @@ -451,6 +454,8 @@ pub(super) enum WebSocketFailurePolicy { pub(super) struct WebSocketRecoveryContext<'a> { pub(super) policy: WebSocketFailurePolicy, pub(super) requirement: TransportRequirement, + /// 本 attempt 已从 metadata 捕获可恢复同一请求的 turn-state 检查点。 + pub(super) replay_checkpoint_available: bool, pub(super) retry_count: u32, pub(super) max_retries: u32, pub(super) request_id: &'a str, @@ -505,10 +510,11 @@ pub(super) fn apply_websocket_recovery_policy( context: WebSocketRecoveryContext<'_>, ) { let send_state = failure.error.send_state(); - if matches!( + let post_send = matches!( send_state, UpstreamSendState::Sent | UpstreamSendState::Ambiguous - ) { + ); + if post_send && !context.replay_checkpoint_available { // payload 已经发送或发送结果不确定时,代理不能用同一业务输入自动重放。 // 当前请求保持终态失败;恢复状态只裁决客户端主动发起的下一次请求。 let session_recovery_transition = match context.session_recovery_probe.take() { @@ -555,6 +561,11 @@ pub(super) fn apply_websocket_recovery_policy( // 传输恢复必须保持原账号可调度;最终 HTTP attempt 若仍失败,再按真实 HTTP // 结果更新账号健康度,避免中间 WS 错误把同账号钉选提前冷却掉。 failure.account_failure = None; + if post_send { + // Core 对 Ambiguous 默认仍然 fail-closed;只有同一 attempt 的 session + // checkpoint 与这个 Provider proof 同时存在时,才允许同账号传输恢复。 + failure.error.set_replay_safe(); + } let fallback_now = context.policy == WebSocketFailurePolicy::ImmediateFallback || context.retry_count >= context.max_retries; if !fallback_now { @@ -583,6 +594,7 @@ pub(super) fn apply_websocket_recovery_policy( websocket_max_retries = context.max_retries, websocket_retry_delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX), transport_requirement = context.requirement.as_str(), + replay_checkpoint_available = context.replay_checkpoint_available, "OpenAI upstream WebSocket failed; retrying the same account" ); return; @@ -621,6 +633,7 @@ pub(super) fn apply_websocket_recovery_policy( WebSocketFailurePolicy::Budgeted => "retry_budget_exhausted", WebSocketFailurePolicy::ImmediateFallback => "upgrade_required", }, + replay_checkpoint_available = context.replay_checkpoint_available, session_affinity_present = context.session_affinity_key.is_some(), session_affinity_key_hash = context.session_affinity_key_hash.unwrap_or(""), session_transport_action = @@ -673,14 +686,31 @@ pub(super) fn continuation_replay_required_error(reason: &'static str) -> Provid pub(super) fn map_stream_error(error: CodexClientError) -> MappedProviderFailure { let allows_pre_delivery_retry = stream_transport_allows_pre_delivery_retry(&error); let websocket_failure = error.transport() == Some(CodexBackendTransport::WebSocket); + let websocket_checkpoint_retryable = websocket_metadata_close_can_resume(&error); let mut failure = map_client_error(error, UpstreamSendState::Sent, false); failure.websocket_transport_retryable = allows_pre_delivery_retry && websocket_failure; + failure.websocket_checkpoint_retryable = + allows_pre_delivery_retry && websocket_checkpoint_retryable; if allows_pre_delivery_retry && !websocket_failure { failure.error = failure.error.with_pre_delivery_retry(); } failure } +fn websocket_metadata_close_can_resume(error: &CodexClientError) -> bool { + let CodexClientError::WebSocket(error) = error else { + return false; + }; + let Some(close) = error.close_before_terminal() else { + return false; + }; + close.code() == Some(1000) + && matches!( + close.last_event_type(), + Some("codex.response.metadata" | "response.metadata") + ) +} + pub(super) fn stream_transport_allows_pre_delivery_retry(error: &CodexClientError) -> bool { match error { CodexClientError::Http(_) @@ -1068,6 +1098,7 @@ pub(super) fn map_upstream_failure( MappedProviderFailure { error, websocket_transport_retryable: false, + websocket_checkpoint_retryable: false, account_failure: account_failure( category, failure.retry_after_seconds, diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 86acc684f..086609610 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -321,12 +321,14 @@ impl Provider for CodexProvider { identity.prepare_local_conversation(&mut upstream_request); } if let Some(previous_session) = previous_session.as_ref() { - upstream_request.turn_state = if same_client_turn( - previous_session.client_turn_id.as_deref(), - upstream_request.client_turn_id.as_deref(), - ) { + upstream_request.turn_state = if previous_session.retry_checkpoint + || same_client_turn( + previous_session.client_turn_id.as_deref(), + upstream_request.client_turn_id.as_deref(), + ) { // 客户端可能携带 metadata-close 失败后才拿到的新状态;同一 - // turn 内显式回传值比最后一次成功响应保存的状态更新。 + // turn 内显式回传值比最后一次成功响应保存的状态更新。同请求 + // 内部恢复由一次性 retry checkpoint 明确证明,不依赖客户端 turn ID。 upstream_request .turn_state .take() diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index fb13fa8b6..ecb2beeff 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -665,6 +665,7 @@ async fn capture_turn_state_request( previous_turn_id: Option<&str>, current_turn_id: Option<&str>, client_turn_state: Option<&str>, + retry_checkpoint: bool, ) -> wiremock::Request { let account_id = "acct_session_affinity"; let store = Arc::new(MemoryAccountStore::default()); @@ -689,6 +690,9 @@ async fn capture_turn_state_request( if let Some(previous_turn_id) = previous_turn_id { session_state.insert("client_turn_id".to_owned(), json!(previous_turn_id)); } + if retry_checkpoint { + session_state.insert("retry_checkpoint".to_owned(), json!(true)); + } let mut protocol_context = Map::from_iter([("use_websocket".to_owned(), json!(false))]); if let Some(current_turn_id) = current_turn_id { protocol_context.insert("turn_id".to_owned(), json!(current_turn_id)); @@ -2521,7 +2525,7 @@ async fn websocket_turn_state_metadata_is_exposed_through_response_observation() } #[tokio::test] -async fn websocket_turn_state_metadata_survives_close_1000_without_changing_failure_semantics() { +async fn websocket_turn_state_metadata_close_requests_checkpointed_retry() { const ACCOUNT_ID: &str = "acct_websocket_metadata_close"; const TURN_STATE: &str = "turn-state-before-close"; @@ -2579,6 +2583,7 @@ async fn websocket_turn_state_metadata_survives_close_1000_without_changing_fail .await .expect("prepare WebSocket provider stream"); let mut observed_turn_state = false; + let mut replay_checkpoint = None; let error = loop { match stream.next().await { Some(Ok(event)) => { @@ -2588,6 +2593,9 @@ async fn websocket_turn_state_metadata_survives_close_1000_without_changing_fail && header.value().as_ref() == TURN_STATE.as_bytes() }); } + if let Some(update) = event.session_update() { + replay_checkpoint = Some(update.clone()); + } } Some(Err(error)) => break error, None => panic!("metadata close must surface a provider error"), @@ -2596,8 +2604,22 @@ async fn websocket_turn_state_metadata_survives_close_1000_without_changing_fail server.await.expect("WebSocket server"); assert!(observed_turn_state); + let replay_checkpoint = replay_checkpoint.expect("metadata close checkpoint"); + assert_eq!( + replay_checkpoint.payload().get("turn_state"), + Some(&json!(TURN_STATE)) + ); + assert_eq!( + replay_checkpoint.payload().get("retry_checkpoint"), + Some(&json!(true)) + ); assert_eq!(error.send_state(), UpstreamSendState::Ambiguous); - assert_eq!(error.pre_delivery_retry(), None); + assert!(error.replay_is_safe()); + assert!(matches!( + error.pre_delivery_retry(), + Some(PreDeliveryRetry::SameAccountTransportRetry { retry_index, .. }) + if retry_index.get() == 1 + )); assert_eq!( error.upstream_code().map(|code| code.as_str()), Some("websocket_close_1000") @@ -3036,6 +3058,7 @@ async fn matching_turn_id_should_restore_previous_turn_state() { Some("turn-same"), Some("turn-same"), None, + false, ) .await; @@ -3045,6 +3068,17 @@ async fn matching_turn_id_should_restore_previous_turn_state() { ); } +#[tokio::test] +async fn retry_checkpoint_should_restore_turn_state_without_client_turn_id() { + let request = + capture_turn_state_request("req_restore_retry_checkpoint", None, None, None, true).await; + + assert_eq!( + captured_header_values(&request, "x-codex-turn-state"), + vec![b"previous-turn-state".to_vec()] + ); +} + #[tokio::test] async fn matching_turn_id_should_prefer_an_explicit_client_echo_over_saved_provider_state() { let request = capture_turn_state_request( @@ -3052,6 +3086,7 @@ async fn matching_turn_id_should_prefer_an_explicit_client_echo_over_saved_provi Some("turn-same"), Some("turn-same"), Some("client-turn-state"), + false, ) .await; @@ -3083,6 +3118,7 @@ async fn new_or_unidentified_turn_should_not_restore_previous_turn_state() { previous_turn_id, current_turn_id, client_turn_state, + false, ) .await; assert!(captured_header_values(&request, "x-codex-turn-state").is_empty());