From 6b362a5029ae4949049c090ebd548c14a3ad0ffe Mon Sep 17 00:00:00 2001 From: Zyy <0x@ainz.cc> Date: Wed, 4 Nov 2026 19:30:22 +0800 Subject: [PATCH] fix(accounts): keep failed recovery probes from thawing accounts (#605) Carry the freeze recovery target into provider account selection so probe responses cannot update credentials or clear the active protection. Let the recovery worker postpone or release the same freeze after observing the upstream result. Co-authored-by: Zyy Co-authored-by: Codex --- .../gateway-admin/src/freeze_recovery.rs | 3 +- .../src/use_case/accounts/mod.rs | 10 +-- .../tests/use_case/accounts/mod.rs | 13 +++- .../tests/use_case/freeze_recovery.rs | 5 ++ .../src/admin/accounts/handlers.rs | 7 +- .../gateway-core/src/account/selection.rs | 13 +++- .../gateway-core/src/engine/coordinator.rs | 17 ++--- .../src/engine/execution/service.rs | 6 +- backend/crates/gateway-core/src/engine/mod.rs | 28 ++++---- .../crates/gateway-core/src/engine/probe.rs | 32 ++++++++- .../gateway-core/tests/engine/coordinator.rs | 14 +++- .../tests/engine/execution/mod.rs | 19 ++++-- .../openai/src/credential/selector.rs | 24 ++++--- .../openai/tests/provider/contract/cookies.rs | 65 +++++++++++++++++++ .../openai/tests/provider/contract/mod.rs | 14 ++-- .../tests/provider/contract/turn_state.rs | 5 +- .../providers/xai/src/credential/selector.rs | 12 ++-- .../crates/providers/xai/src/provider/mod.rs | 9 +-- .../xai/tests/provider/contract/mod.rs | 7 +- docs/architecture.md | 3 +- 20 files changed, 227 insertions(+), 79 deletions(-) diff --git a/backend/crates/gateway-admin/src/freeze_recovery.rs b/backend/crates/gateway-admin/src/freeze_recovery.rs index 00bdecfd2..c4acdd455 100644 --- a/backend/crates/gateway-admin/src/freeze_recovery.rs +++ b/backend/crates/gateway-admin/src/freeze_recovery.rs @@ -7,6 +7,7 @@ use std::time::{Duration, SystemTime}; use chrono::{DateTime, Utc}; use futures::StreamExt as _; use gateway_core::account::{AccountErrorProtection, ProviderAccountId}; +use gateway_core::engine::probe::AccountProbeTarget; use gateway_core::routing::UpstreamModelId; use gateway_core::task::{ ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerId, WorkerKind, @@ -161,7 +162,7 @@ impl FreezeRecoveryTask { let probe_succeeded = match self .deps .accounts - .test_connection(account.clone(), model) + .test_connection(AccountProbeTarget::FreezeRecovery(account.clone()), model) .await { Ok(events) => drain_probe(events).await, diff --git a/backend/crates/gateway-admin/src/use_case/accounts/mod.rs b/backend/crates/gateway-admin/src/use_case/accounts/mod.rs index cdc68cf2c..e2c6a9e88 100644 --- a/backend/crates/gateway-admin/src/use_case/accounts/mod.rs +++ b/backend/crates/gateway-admin/src/use_case/accounts/mod.rs @@ -12,7 +12,7 @@ use chrono::{Duration, Utc}; use futures::StreamExt as _; use gateway_core::{ account::ProviderAccountId, - engine::probe::{AccountProbe, AccountProbeRequest}, + engine::probe::{AccountProbe, AccountProbeRequest, AccountProbeTarget}, routing::{ProviderKind, UpstreamModelId}, runtime::SnapshotControl, }; @@ -178,7 +178,7 @@ pub trait AccountsService: Send + Sync { async fn test_connection( &self, - account_id: ProviderAccountId, + target: AccountProbeTarget, upstream_model: UpstreamModelId, ) -> Result; } @@ -1131,10 +1131,10 @@ impl AccountsService for DefaultAccountsService { async fn test_connection( &self, - account_id: ProviderAccountId, + target: AccountProbeTarget, upstream_model: UpstreamModelId, ) -> Result { - let (stored, provider) = self.provider_for_account(&account_id).await?; + let (stored, provider) = self.provider_for_account(target.account_id()).await?; let account = stored.account; let model = upstream_model.as_str().to_owned(); let operation = provider @@ -1157,7 +1157,7 @@ impl AccountsService for DefaultAccountsService { let result = probe .probe( AccountProbeRequest { - account_id, + target, provider_kind: account.provider_kind, upstream_model, operation, diff --git a/backend/crates/gateway-admin/tests/use_case/accounts/mod.rs b/backend/crates/gateway-admin/tests/use_case/accounts/mod.rs index dfdc3dc6e..8bd27b700 100644 --- a/backend/crates/gateway-admin/tests/use_case/accounts/mod.rs +++ b/backend/crates/gateway-admin/tests/use_case/accounts/mod.rs @@ -2,6 +2,7 @@ mod reset_credit_automation; +use gateway_core::engine::probe::AccountProbeTarget; use std::{ collections::BTreeMap, sync::{Arc, Mutex}, @@ -1335,7 +1336,9 @@ async fn connection_test_should_probe_unavailable_account() { let events = services .accounts() .test_connection( - ProviderAccountId::new("acct_test").expect("account ID"), + AccountProbeTarget::ConnectionTest( + ProviderAccountId::new("acct_test").expect("account ID"), + ), gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"), ) .await @@ -1372,7 +1375,9 @@ async fn connection_test_rate_limited_probe_returns_provider_failure() { let events = services .accounts() .test_connection( - ProviderAccountId::new("acct_test").expect("account ID"), + AccountProbeTarget::ConnectionTest( + ProviderAccountId::new("acct_test").expect("account ID"), + ), gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"), ) .await @@ -1405,7 +1410,9 @@ async fn connection_test_should_preserve_disabled_account_status() { let events = services .accounts() .test_connection( - ProviderAccountId::new("acct_test").expect("account ID"), + AccountProbeTarget::ConnectionTest( + ProviderAccountId::new("acct_test").expect("account ID"), + ), gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"), ) .await diff --git a/backend/crates/gateway-admin/tests/use_case/freeze_recovery.rs b/backend/crates/gateway-admin/tests/use_case/freeze_recovery.rs index 760014f1c..a5d7a46d4 100644 --- a/backend/crates/gateway-admin/tests/use_case/freeze_recovery.rs +++ b/backend/crates/gateway-admin/tests/use_case/freeze_recovery.rs @@ -1,5 +1,6 @@ //! 账号保护恢复编排：到期探测、快照参数与失败顺延 +use gateway_core::engine::probe::AccountProbeTarget; use std::collections::BTreeMap; use std::sync::{Arc, Mutex}; @@ -242,6 +243,10 @@ impl AccountProbe for FailingProbe { request: AccountProbeRequest, _: Option>, ) -> futures::future::BoxFuture<'_, Result> { + assert!(matches!( + request.target, + AccountProbeTarget::FreezeRecovery(_) + )); assert_eq!(request.upstream_model.as_str(), "snapshot-model"); Box::pin(async { Err(AccountProbeError::new( diff --git a/backend/crates/gateway-api/src/admin/accounts/handlers.rs b/backend/crates/gateway-api/src/admin/accounts/handlers.rs index 105f8ea66..556817b48 100644 --- a/backend/crates/gateway-api/src/admin/accounts/handlers.rs +++ b/backend/crates/gateway-api/src/admin/accounts/handlers.rs @@ -1,5 +1,7 @@ //! 账号管理路由与 HTTP handler 编排 +use gateway_core::engine::probe::AccountProbeTarget; + use crate::auth::SessionState; use super::*; @@ -626,7 +628,10 @@ where let stream = state .admin_services() .accounts() - .test_connection(account_id, upstream_model) + .test_connection( + AccountProbeTarget::ConnectionTest(account_id), + upstream_model, + ) .await .map_err(map_service_error)? .map(move |event| { diff --git a/backend/crates/gateway-core/src/account/selection.rs b/backend/crates/gateway-core/src/account/selection.rs index 93ad8f4c9..318dacce8 100644 --- a/backend/crates/gateway-core/src/account/selection.rs +++ b/backend/crates/gateway-core/src/account/selection.rs @@ -551,12 +551,23 @@ pub enum AccountEligibilityPolicy { #[default] Enforce, BypassForDiagnostic, + FreezeRecovery, } impl AccountEligibilityPolicy { #[must_use] pub const fn bypasses_local_eligibility(self) -> bool { - matches!(self, Self::BypassForDiagnostic) + !matches!(self, Self::Enforce) + } + + /// 恢复探测只验证上游结果，不能通过回写 Cookie 或刷新凭据提前解除冻结 + #[must_use] + pub const fn allows_account_state_mutation(self, enabled: bool) -> bool { + match self { + Self::Enforce => true, + Self::BypassForDiagnostic => enabled, + Self::FreezeRecovery => false, + } } } diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index 65f5d83a6..544cc31ad 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -18,6 +18,7 @@ use crate::engine::continuation::{ ContinuationBinding, NativeContinuationPin, NativeContinuationScope, PreviousResponseId, }; +use crate::engine::probe::AccountProbeTarget; use crate::engine::provider::{Provider, ProviderCallMetadata, ProviderRequest, ProviderStream}; use crate::engine::{ AccountAttemptContext, AttemptContext, AttemptRecord, AttemptTransport, AttemptTrigger, @@ -49,7 +50,7 @@ pub struct AttemptCoordinator { #[derive(Debug, Clone)] enum AccountSelection { Scheduled(Option), - Diagnostic(crate::account::ProviderAccountId), + Diagnostic(AccountProbeTarget), } pub(super) struct CoordinationExtensions { @@ -153,7 +154,7 @@ impl AccountSelection { fn required_account(&self) -> Option<&crate::account::ProviderAccountId> { match self { Self::Scheduled(account) => account.as_ref(), - Self::Diagnostic(account) => Some(account), + Self::Diagnostic(target) => Some(target.account_id()), } } @@ -226,7 +227,7 @@ where request: NewModelRequest, operation: Operation, plan: RoutingPlan, - required_account: crate::account::ProviderAccountId, + target: AccountProbeTarget, continuation: Option, cancellation: CancellationToken, ) -> Result, EngineError> { @@ -234,7 +235,7 @@ where request, operation, plan, - AccountSelection::Diagnostic(required_account), + AccountSelection::Diagnostic(target), CoordinationExtensions::new(continuation, None), cancellation, ) @@ -1086,8 +1087,8 @@ where (Some(recovery.account), recovery.transport) } else { match &self.account_selection { - AccountSelection::Diagnostic(account) => { - (Some(account.clone()), AttemptTransport::Default) + AccountSelection::Diagnostic(target) => { + (Some(target.account_id().clone()), AttemptTransport::Default) } AccountSelection::Scheduled(_) => ( self.recovery_account @@ -1100,9 +1101,9 @@ where } }; let account_context = match &self.account_selection { - AccountSelection::Diagnostic(account) => AccountAttemptContext::diagnostic( + AccountSelection::Diagnostic(target) => AccountAttemptContext::diagnostic( self.excluded_accounts.clone(), - account.clone(), + target.clone(), self.account_state_owner.clone(), ), AccountSelection::Scheduled(_) => AccountAttemptContext::new( diff --git a/backend/crates/gateway-core/src/engine/execution/service.rs b/backend/crates/gateway-core/src/engine/execution/service.rs index 31d95d78d..00552d3a3 100644 --- a/backend/crates/gateway-core/src/engine/execution/service.rs +++ b/backend/crates/gateway-core/src/engine/execution/service.rs @@ -1334,14 +1334,14 @@ impl DefaultExecutionService { frozen: Option>, ) -> Result { let AccountProbeRequest { - account_id, + target, provider_kind, upstream_model, operation, } = request; let observed = ProbeObservation { provider_kind: provider_kind.clone(), - account_id: account_id.clone(), + account_id: target.account_id().clone(), upstream_model: upstream_model.clone(), }; let snapshot = frozen @@ -1402,7 +1402,7 @@ impl DefaultExecutionService { new_request, operation, plan, - account_id, + target, None, CancellationToken::new(), ) diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index e44271daf..732d5d497 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -35,8 +35,8 @@ use async_trait::async_trait; use thiserror::Error; use crate::account::{ - AccountCandidate, AccountSelection, AccountSelectionContext, AccountSelectionPolicy, FastMode, - ProviderAccountId, + AccountCandidate, AccountEligibilityPolicy, AccountSelection, AccountSelectionContext, + AccountSelectionPolicy, FastMode, ProviderAccountId, }; use crate::engine::continuation::{ContinuationBinding, NativeContinuationPin}; use crate::error::{ @@ -128,7 +128,7 @@ pub struct AccountAttemptContext { required_account: Option, state_owner: Option, credential_recovery_attempted: bool, - diagnostic_required_account: bool, + eligibility: AccountEligibilityPolicy, account_scope: Option>, } @@ -144,7 +144,7 @@ impl AccountAttemptContext { required_account, state_owner, credential_recovery_attempted: false, - diagnostic_required_account: false, + eligibility: AccountEligibilityPolicy::Enforce, account_scope: None, } } @@ -153,17 +153,18 @@ impl AccountAttemptContext { /// /// 诊断只绕过本地的被动可用性投影，不能换号，并继续受账号租约保护 #[must_use] - pub const fn diagnostic( + pub fn diagnostic( excluded_accounts: BTreeSet, - required_account: ProviderAccountId, + target: probe::AccountProbeTarget, state_owner: Option, ) -> Self { + let (required_account, eligibility) = target.into_selection(); Self { excluded_accounts, required_account: Some(required_account), state_owner, credential_recovery_attempted: false, - diagnostic_required_account: true, + eligibility, account_scope: None, } } @@ -205,12 +206,6 @@ impl AccountAttemptContext { self.credential_recovery_attempted } - /// 此尝试是否要由真实上游而非本地投影确认固定账号的可用性 - #[must_use] - pub const fn is_diagnostic_required_account(&self) -> bool { - self.diagnostic_required_account - } - #[must_use] pub const fn account_scope(&self) -> Option<&Arc> { self.account_scope.as_ref() @@ -921,7 +916,12 @@ impl AttemptContext { /// 管理端固定账号诊断会跳过本地可用性投影，保留全部租约约束 #[must_use] pub const fn is_diagnostic_required_account(&self) -> bool { - self.account.is_diagnostic_required_account() + self.account_eligibility().bypasses_local_eligibility() + } + + #[must_use] + pub const fn account_eligibility(&self) -> AccountEligibilityPolicy { + self.account.eligibility } /// 当前路由候选的冻结账号范围；管理端诊断为 `None` diff --git a/backend/crates/gateway-core/src/engine/probe.rs b/backend/crates/gateway-core/src/engine/probe.rs index 068e7ecf3..0330cdddc 100644 --- a/backend/crates/gateway-core/src/engine/probe.rs +++ b/backend/crates/gateway-core/src/engine/probe.rs @@ -8,11 +8,39 @@ use futures::future::BoxFuture; use crate::error::{ClientVisibleUpstreamResponse, GatewayError, GatewayErrorKind}; use crate::identity::ProviderKind; use crate::routing::UpstreamModelId; -use crate::{account::ProviderAccountId, operation::Operation, upstream::UpstreamSendState}; +use crate::{ + account::{AccountEligibilityPolicy, ProviderAccountId}, + operation::Operation, + upstream::UpstreamSendState, +}; + +/// 探测始终固定账号，恢复探测的账号状态由解冻任务持有 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AccountProbeTarget { + ConnectionTest(ProviderAccountId), + FreezeRecovery(ProviderAccountId), +} + +impl AccountProbeTarget { + pub fn account_id(&self) -> &ProviderAccountId { + match self { + Self::ConnectionTest(account) | Self::FreezeRecovery(account) => account, + } + } + + pub(super) fn into_selection(self) -> (ProviderAccountId, AccountEligibilityPolicy) { + match self { + Self::ConnectionTest(account) => { + (account, AccountEligibilityPolicy::BypassForDiagnostic) + } + Self::FreezeRecovery(account) => (account, AccountEligibilityPolicy::FreezeRecovery), + } + } +} #[derive(Debug, Clone, PartialEq)] pub struct AccountProbeRequest { - pub account_id: ProviderAccountId, + pub target: AccountProbeTarget, pub provider_kind: ProviderKind, pub upstream_model: UpstreamModelId, pub operation: Operation, diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index ea1d829f0..cb04fcee0 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -1,5 +1,6 @@ //! 单行 `model_requests`、账号重试与下游提交屏障测试 +use gateway_core::engine::probe::AccountProbeTarget; use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::num::NonZeroU32; use std::sync::{ @@ -2753,7 +2754,7 @@ fn continuation_history_failure_without_replay_proof_is_delivered_without_hidden } #[test] -fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() { +fn recovery_probe_reaches_provider_with_read_only_account() { let operation = generate_operation(); let route_plan = plan(&operation); let (coordinator, _, provider) = coordinator(vec![Script::Stream { @@ -2765,7 +2766,7 @@ fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() { model_request(&operation, SystemTime::now() + Duration::from_secs(30)), operation, route_plan, - required.clone(), + AccountProbeTarget::FreezeRecovery(required.clone()), None, CancellationToken::new(), )) @@ -2777,6 +2778,11 @@ fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() { assert_eq!(contexts.len(), 1); assert_eq!(contexts[0].required_account(), Some(&required)); assert!(contexts[0].is_diagnostic_required_account()); + assert!( + !contexts[0] + .account_eligibility() + .allows_account_state_mutation(true) + ); } #[test] @@ -7397,7 +7403,9 @@ fn error_protection_tracks_real_retry_attempts_and_skips_diagnostic_runs() { request, operation, route_plan, - ProviderAccountId::new("acct_first").unwrap(), + AccountProbeTarget::ConnectionTest( + ProviderAccountId::new("acct_first").unwrap(), + ), None, CancellationToken::new(), ) diff --git a/backend/crates/gateway-core/tests/engine/execution/mod.rs b/backend/crates/gateway-core/tests/engine/execution/mod.rs index d899df6ee..1ab5d088c 100644 --- a/backend/crates/gateway-core/tests/engine/execution/mod.rs +++ b/backend/crates/gateway-core/tests/engine/execution/mod.rs @@ -2,6 +2,7 @@ mod identity; mod outbound; +use gateway_core::engine::probe::AccountProbeTarget; use gateway_core::routing::{ ContributedModel, ContributedModelAlias, extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference}, @@ -3135,7 +3136,9 @@ fn account_probe_should_not_write_to_the_persistent_execution_store() { let error = block_on(service.probe( AccountProbeRequest { - account_id: ProviderAccountId::new("acct_probe").expect("account ID"), + target: AccountProbeTarget::ConnectionTest( + ProviderAccountId::new("acct_probe").expect("account ID"), + ), provider_kind: ProviderKind::new("openai").expect("provider kind"), upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"), operation: probe_operation(), @@ -3167,7 +3170,9 @@ fn probe_failures_should_be_observable_without_a_model_request_row() { let error = block_on(service.probe( AccountProbeRequest { - account_id: ProviderAccountId::new("acct_probe").expect("account ID"), + target: AccountProbeTarget::ConnectionTest( + ProviderAccountId::new("acct_probe").expect("account ID"), + ), provider_kind: ProviderKind::new("openai").expect("provider kind"), upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"), operation: probe_operation(), @@ -3213,7 +3218,9 @@ fn provider_local_probe_failure_should_remain_distinct_from_upstream() { let error = block_on(service.probe( AccountProbeRequest { - account_id: ProviderAccountId::new("acct_probe").expect("account ID"), + target: AccountProbeTarget::ConnectionTest( + ProviderAccountId::new("acct_probe").expect("account ID"), + ), provider_kind: ProviderKind::new("openai").expect("provider kind"), upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"), operation: probe_operation(), @@ -3256,7 +3263,7 @@ fn diagnostic_probe_does_not_apply_data_plane_account_model_policy() { let error = block_on(service.probe( AccountProbeRequest { - account_id, + target: AccountProbeTarget::ConnectionTest(account_id), provider_kind: provider, upstream_model: UpstreamModelId::new("gpt-probe").expect("model"), operation: probe_operation(), @@ -3288,7 +3295,9 @@ fn probe_observation_store_failure_preserves_the_provider_error() { let error = block_on(service.probe( AccountProbeRequest { - account_id: ProviderAccountId::new("acct_probe").expect("account ID"), + target: AccountProbeTarget::ConnectionTest( + ProviderAccountId::new("acct_probe").expect("account ID"), + ), provider_kind: ProviderKind::new("openai").expect("provider kind"), upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"), operation: probe_operation(), diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index 4763a7cf7..4417652e5 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -8,11 +8,11 @@ use std::time::{Duration, SystemTime}; use gateway_core::error::ErrorSource; use gateway_core::account::{ - AccountAffinity, AccountCandidate, AccountCapacitySnapshot, AccountEligibilityPolicy, - AccountErrorReason, AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, - AccountSchedulingBlocker, AccountSelectionContext, AccountSelectionPolicy, AccountSelector, - AccountStatus, CredentialState, PreferredAccountSelection, ProviderAccount, ProviderAccountId, - QuotaEvidence, RotationStrategy, + AccountAffinity, AccountCandidate, AccountCapacitySnapshot, AccountErrorReason, + AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker, + AccountSelectionContext, AccountSelectionPolicy, AccountSelector, AccountStatus, + CredentialState, PreferredAccountSelection, ProviderAccount, ProviderAccountId, QuotaEvidence, + RotationStrategy, }; use gateway_core::concurrency::{CapacityWait, ConcurrencyWaitQueue, QueueRejection, WaitPriority}; use gateway_core::engine::{AttemptContext, ContinuationAttempt, policy::AccountPolicyError}; @@ -510,7 +510,8 @@ impl CodexCredentialSelector { ), _ => None, }; - let diagnostic = request.attempt.is_diagnostic_required_account(); + let eligibility = request.attempt.account_eligibility(); + let diagnostic = eligibility.bypasses_local_eligibility(); let binding_key = request.session_affinity_key.filter(|_| !diagnostic); let affinity_mode = request .attempt @@ -861,11 +862,7 @@ impl CodexCredentialSelector { preferred_account: preferred.clone(), preferred_account_overrides_weight: true, round_robin_cursor, - eligibility: if diagnostic { - AccountEligibilityPolicy::BypassForDiagnostic - } else { - AccountEligibilityPolicy::Enforce - }, + eligibility, account_scope: request.attempt.account_scope().cloned(), reserved_concurrency, }; @@ -1062,7 +1059,8 @@ impl CodexCredentialSelector { busy_accounts.insert(account.id().clone()); continue; } - let allows_account_state_mutation = !diagnostic || account.enabled(); + let allows_account_state_mutation = + eligibility.allows_account_state_mutation(account.enabled()); match self .leases .try_acquire(ProviderLeaseRequest::Scheduling( @@ -1814,7 +1812,7 @@ impl CodexCredentialLease { self.cyber_policy_scope.as_ref() } - /// 禁用账号的管理端诊断必须只返回真实上游结果，不能回写账号侧状态 + /// 租约保留本次选号的回写权限，覆盖成功、失败与流式反馈 #[must_use] pub(crate) const fn allows_account_state_mutation(&self) -> bool { self.allows_account_state_mutation diff --git a/backend/crates/providers/openai/tests/provider/contract/cookies.rs b/backend/crates/providers/openai/tests/provider/contract/cookies.rs index 33c74f08f..19b3dbf21 100644 --- a/backend/crates/providers/openai/tests/provider/contract/cookies.rs +++ b/backend/crates/providers/openai/tests/provider/contract/cookies.rs @@ -52,6 +52,71 @@ fn cookie_provider(store: &Arc, base_url: String) -> Arc { + completed |= event + .canonical_facts() + .iter() + .any(|fact| matches!(fact, GatewayEvent::Completed(_))) + } + Err(error) => { + assert_eq!(status, 400, "expected a completed probe: {error:?}"); + rejected = error.upstream_status(); + break; + } + } + } + assert_eq!(completed, status == 200); + assert_eq!(rejected, (status == 400).then_some(400)); + assert_eq!( + store.account(account_id).unwrap().revision(), + original.revision() + ); + } +} + #[tokio::test] async fn outbound_redirected_response_cookies_keep_the_actual_origin_on_success_and_failure() { for status in [200, 400] { diff --git a/backend/crates/providers/openai/tests/provider/contract/mod.rs b/backend/crates/providers/openai/tests/provider/contract/mod.rs index c9ee5b3d9..a9f64e6a4 100644 --- a/backend/crates/providers/openai/tests/provider/contract/mod.rs +++ b/backend/crates/providers/openai/tests/provider/contract/mod.rs @@ -46,6 +46,7 @@ use gateway_core::engine::middleware::{ FrozenMiddlewarePlan, MiddlewareContext, MiddlewareError, MiddlewareHeader, MiddlewareMount, MiddlewareNext, MiddlewarePlan, MiddlewareRequest, MiddlewareResponse, }; +use gateway_core::engine::probe::AccountProbeTarget; use gateway_core::engine::provider::{Provider as _, ProviderRequest}; use gateway_core::engine::{ AccountAttemptContext, AttemptContext, AttemptTransport, ContinuationAttempt, ModelRequestId, @@ -1689,6 +1690,13 @@ fn fallback_transport_context(request_id: &str) -> AttemptContext { } fn diagnostic_context(request_id: &str, account_id: &str) -> AttemptContext { + probe_context( + request_id, + AccountProbeTarget::ConnectionTest(ProviderAccountId::new(account_id).unwrap()), + ) +} + +fn probe_context(request_id: &str, target: AccountProbeTarget) -> AttemptContext { AttemptContext::new( RequestAttemptContext::new( ModelRequestId::new(request_id).expect("request id"), @@ -1698,11 +1706,7 @@ fn diagnostic_context(request_id: &str, account_id: &str) -> AttemptContext { NonZeroU32::new(1).expect("attempt"), SystemTime::now() + Duration::from_secs(30), account_policy(), - AccountAttemptContext::diagnostic( - BTreeSet::new(), - ProviderAccountId::new(account_id).expect("account id"), - None, - ), + AccountAttemptContext::diagnostic(BTreeSet::new(), target, None), None, CancellationToken::new(), ) diff --git a/backend/crates/providers/openai/tests/provider/contract/turn_state.rs b/backend/crates/providers/openai/tests/provider/contract/turn_state.rs index 8692b2a95..0fea9bcf2 100644 --- a/backend/crates/providers/openai/tests/provider/contract/turn_state.rs +++ b/backend/crates/providers/openai/tests/provider/contract/turn_state.rs @@ -1,6 +1,7 @@ //! 验证同轮全量请求回传的 token 归属、失败路径与账号状态隔离 use super::*; +use gateway_core::engine::probe::AccountProbeTarget; const TOKEN: &str = "upstream-turn-state"; @@ -293,7 +294,9 @@ async fn full_request_turn_state_requires_matching_client_account_thread_turn_an account_policy(), AccountAttemptContext::diagnostic( BTreeSet::new(), - ProviderAccountId::new("acct_session_affinity").unwrap(), + AccountProbeTarget::ConnectionTest( + ProviderAccountId::new("acct_session_affinity").unwrap(), + ), None, ), None, diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index a9b116670..13a223552 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -6,9 +6,9 @@ use std::time::{Duration, SystemTime}; use chrono::Utc; use gateway_core::account::{ - AccountCandidate, AccountCapacitySnapshot, AccountEligibilityPolicy, AccountFeedbackStats, - AccountSelectionContext, AccountSelector, CredentialState, ProviderAccount, ProviderAccountId, - QuotaAccessState, QuotaEvidence, QuotaState, + AccountCandidate, AccountCapacitySnapshot, AccountFeedbackStats, AccountSelectionContext, + AccountSelector, CredentialState, ProviderAccount, ProviderAccountId, QuotaAccessState, + QuotaEvidence, QuotaState, }; use gateway_core::concurrency::{CapacityWait, ConcurrencyWaitQueue}; use gateway_core::engine::policy::AccountPolicyError; @@ -101,7 +101,7 @@ impl GrokAccountSessionSelector { request.concurrency_wait_budget(), ); 'refresh: loop { - let diagnostic = request.eligibility() == AccountEligibilityPolicy::BypassForDiagnostic; + let diagnostic = request.eligibility().bypasses_local_eligibility(); // store 侧常规调度列表不包含停用账号；管理端诊断要对固定账号执行真实上游 // 验证，这里把不在列表里的 required 账号显式补回候选 let mut accounts = @@ -315,7 +315,9 @@ impl GrokAccountSessionSelector { let selected = selection.candidate(); let selected_id = selected.account.id().clone(); let selected_revision = selected.account.revision(); - let allows_account_state_mutation = !diagnostic || selected.account.enabled(); + let allows_account_state_mutation = request + .eligibility() + .allows_account_state_mutation(selected.account.enabled()); let lease = self .scheduling .try_acquire(ProviderLeaseRequest::Scheduling( diff --git a/backend/crates/providers/xai/src/provider/mod.rs b/backend/crates/providers/xai/src/provider/mod.rs index 546a7f001..88bfc2c4e 100644 --- a/backend/crates/providers/xai/src/provider/mod.rs +++ b/backend/crates/providers/xai/src/provider/mod.rs @@ -7,8 +7,7 @@ use std::time::{Duration, Instant, SystemTime}; use async_trait::async_trait; use futures::{StreamExt, future::BoxFuture}; use gateway_core::account::{ - AccountEligibilityPolicy, AccountFeedbackStats, ProviderAccount, ProviderAccountId, - ProviderAccountStore, + AccountFeedbackStats, ProviderAccount, ProviderAccountId, ProviderAccountStore, }; use gateway_core::engine::continuation::ContinuationBinding; use gateway_core::engine::middleware::MiddlewareHeader; @@ -725,11 +724,7 @@ async fn select_grok_session( context.request_policy_context().cloned(), context.attempt_index(), ) - .with_eligibility_policy(if context.is_diagnostic_required_account() { - AccountEligibilityPolicy::BypassForDiagnostic - } else { - AccountEligibilityPolicy::Enforce - }) + .with_eligibility_policy(context.account_eligibility()) .with_affinity(affinity); selection.upstream_context = Some(context.clone()); selection.adapter_required = adapter_required; diff --git a/backend/crates/providers/xai/tests/provider/contract/mod.rs b/backend/crates/providers/xai/tests/provider/contract/mod.rs index d44e20763..b1be3db38 100644 --- a/backend/crates/providers/xai/tests/provider/contract/mod.rs +++ b/backend/crates/providers/xai/tests/provider/contract/mod.rs @@ -6,6 +6,7 @@ mod scheduling; mod selection_retry; mod timing; +use gateway_core::engine::probe::AccountProbeTarget; use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::num::NonZeroU32; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -1978,7 +1979,11 @@ fn diagnostic_context(request_id: &str, account_suffix: &str) -> AttemptContext NonZeroU32::new(1).expect("attempt"), SystemTime::now() + Duration::from_secs(30), selection_policy(), - AccountAttemptContext::diagnostic(BTreeSet::new(), account_id(account_suffix), None), + AccountAttemptContext::diagnostic( + BTreeSet::new(), + AccountProbeTarget::ConnectionTest(account_id(account_suffix)), + None, + ), None, CancellationToken::new(), ) diff --git a/docs/architecture.md b/docs/architecture.md index 50d7b71d7..c7cc9b0b9 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -1205,7 +1205,8 @@ Redis 按账号、凭据与策略代次隔离计数，原子完成递增、有 反馈持有在请求生命周期内，单次存储等待有界，失败保留原始客户端结果，不启动脱离请求的后台写入 冻结保留规则标识和名称以及恢复时长、模型快照，不保存原始错误正文。需要探测的冻结到期后仍阻止调度， -Admin worker 复用连接测试，按凭据及冻结代次比较后解除或顺延；旧探测不能覆盖新冻结或管理员操作。 +Admin worker 复用固定账号连接测试并指定恢复用途。Provider 持有账号租约，只读取账号状态，避免响应 Cookie 等反馈推进凭据版本。 +Worker 按探测结果、凭据版本及冻结代次解除或顺延保护；旧探测不能覆盖新冻结或管理员操作。 关闭保护或探测后，已有冻结在当前截止时间后解除；普通成功不解除已有保护。匹配与资源限制见[错误保护 API](api.md#账号错误保护) 周期任务的 Redis lease 只保证单周期互斥，不构成多副本 leader 选举。备份 daemon 依赖单副本部署边界，