mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
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 <zyycn@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
Zyy
Codex
@@ -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,
|
||||
|
||||
@@ -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<AccountConnectionTestEventStream, AdminError>;
|
||||
}
|
||||
@@ -1131,10 +1131,10 @@ impl AccountsService for DefaultAccountsService {
|
||||
|
||||
async fn test_connection(
|
||||
&self,
|
||||
account_id: ProviderAccountId,
|
||||
target: AccountProbeTarget,
|
||||
upstream_model: UpstreamModelId,
|
||||
) -> Result<AccountConnectionTestEventStream, AdminError> {
|
||||
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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<Arc<gateway_core::routing::RuntimeSnapshot>>,
|
||||
) -> futures::future::BoxFuture<'_, Result<AccountProbeResult, AccountProbeError>> {
|
||||
assert!(matches!(
|
||||
request.target,
|
||||
AccountProbeTarget::FreezeRecovery(_)
|
||||
));
|
||||
assert_eq!(request.upstream_model.as_str(), "snapshot-model");
|
||||
Box::pin(async {
|
||||
Err(AccountProbeError::new(
|
||||
|
||||
@@ -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| {
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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<S: ?Sized> {
|
||||
#[derive(Debug, Clone)]
|
||||
enum AccountSelection {
|
||||
Scheduled(Option<crate::account::ProviderAccountId>),
|
||||
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<ContinuationBinding>,
|
||||
cancellation: CancellationToken,
|
||||
) -> Result<ResponseExecutionSession<S>, 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(
|
||||
|
||||
@@ -1334,14 +1334,14 @@ impl DefaultExecutionService {
|
||||
frozen: Option<Arc<crate::routing::RuntimeSnapshot>>,
|
||||
) -> Result<AccountProbeResult, AccountProbeError> {
|
||||
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(),
|
||||
)
|
||||
|
||||
@@ -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<ProviderAccountId>,
|
||||
state_owner: Option<ProviderAccountStateOwner>,
|
||||
credential_recovery_attempted: bool,
|
||||
diagnostic_required_account: bool,
|
||||
eligibility: AccountEligibilityPolicy,
|
||||
account_scope: Option<Arc<crate::account::scope::FrozenAccountScope>>,
|
||||
}
|
||||
|
||||
@@ -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<ProviderAccountId>,
|
||||
required_account: ProviderAccountId,
|
||||
target: probe::AccountProbeTarget,
|
||||
state_owner: Option<ProviderAccountStateOwner>,
|
||||
) -> 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<crate::account::scope::FrozenAccountScope>> {
|
||||
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`
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
)
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -52,6 +52,71 @@ fn cookie_provider(store: &Arc<MemoryAccountStore>, base_url: String) -> Arc<Cod
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn freeze_recovery_receives_upstream_outcome_without_rotating_response_cookies() {
|
||||
let completed_sse = format!(
|
||||
"event: response.created\ndata: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_scope_capture\",\"model\":\"gpt-5.4\"}}}}\n\n{CAPTURE_COMPLETED_SSE}"
|
||||
);
|
||||
for (status, content_type, body) in [
|
||||
(
|
||||
400,
|
||||
"application/json",
|
||||
r#"{"error":{"message":"invalid model","type":"invalid_request_error"}}"#,
|
||||
),
|
||||
(200, "text/event-stream", completed_sse.as_str()),
|
||||
] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_provider_contract";
|
||||
create_account(&store, account_id).await;
|
||||
let original = store.account(account_id).unwrap();
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(status)
|
||||
.insert_header("content-type", content_type)
|
||||
.insert_header("set-cookie", "cf_clearance=changed; Path=/; Max-Age=3600")
|
||||
.set_body_string(body),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let mut stream = cookie_provider(&store, server.uri())
|
||||
.execute(
|
||||
planned_request("openai", http_generate_operation()),
|
||||
probe_context(
|
||||
"req_freeze_recovery",
|
||||
AccountProbeTarget::FreezeRecovery(original.id().clone()),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut completed = false;
|
||||
let mut rejected = None;
|
||||
while let Some(event) = stream.next().await {
|
||||
match event {
|
||||
Ok(event) => {
|
||||
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] {
|
||||
|
||||
@@ -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(),
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(),
|
||||
)
|
||||
|
||||
@@ -1205,7 +1205,8 @@ Redis 按账号、凭据与策略代次隔离计数，原子完成递增、有
|
||||
反馈持有在请求生命周期内，单次存储等待有界，失败保留原始客户端结果，不启动脱离请求的后台写入
|
||||
|
||||
冻结保留规则标识和名称以及恢复时长、模型快照，不保存原始错误正文。需要探测的冻结到期后仍阻止调度，
|
||||
Admin worker 复用连接测试，按凭据及冻结代次比较后解除或顺延；旧探测不能覆盖新冻结或管理员操作。
|
||||
Admin worker 复用固定账号连接测试并指定恢复用途。Provider 持有账号租约，只读取账号状态，避免响应 Cookie 等反馈推进凭据版本。
|
||||
Worker 按探测结果、凭据版本及冻结代次解除或顺延保护；旧探测不能覆盖新冻结或管理员操作。
|
||||
关闭保护或探测后，已有冻结在当前截止时间后解除；普通成功不解除已有保护。匹配与资源限制见[错误保护 API](api.md#账号错误保护)
|
||||
|
||||
周期任务的 Redis lease 只保证单周期互斥，不构成多副本 leader 选举。备份 daemon 依赖单副本部署边界，
|
||||
|
||||
Reference in New Issue
Block a user