mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
feat(oauth): 完善账号异常识别并在调度/展示层拦截失效 OAuth 密钥
- 新增 aether-admin provider status 模块,统一解析账号状态(禁用/工作区停用等) - 调度器 runtime 增加 oauth_invalid 判定,跳过刷新失败或已撤销的 OAuth 密钥(REQUEST_FAILED 保留可选) - gateway state 在 local oauth 刷新返回 4xx 时持久化失败原因并同步状态快照 - admin pool 列表/详情回填 account 状态与 scheduling 阻塞原因(account_blocked) - 共享 catalog 的 status_snapshot payload 附加 account 字段
This commit is contained in:
@@ -22,6 +22,7 @@ pub(super) struct CandidateRuntimeSelectionSnapshot {
|
||||
pub(super) provider_key_rpm_states: BTreeMap<String, StoredProviderCatalogKey>,
|
||||
provider_quota_blocks_requests: BTreeMap<String, bool>,
|
||||
key_account_quota_exhausted: BTreeMap<String, bool>,
|
||||
key_oauth_invalid: BTreeMap<String, bool>,
|
||||
provider_key_rpm_reset_ats: BTreeMap<String, Option<u64>>,
|
||||
}
|
||||
|
||||
@@ -40,6 +41,8 @@ pub(super) async fn read_candidate_runtime_selection_snapshot(
|
||||
&provider_key_rpm_states,
|
||||
&provider_skip_exhausted_accounts,
|
||||
);
|
||||
let key_oauth_invalid =
|
||||
read_key_oauth_invalid_map(candidates, &provider_key_rpm_states, now_unix_secs);
|
||||
let provider_quota_blocks_requests =
|
||||
read_provider_quota_block_map(state, candidates, now_unix_secs).await?;
|
||||
let provider_key_rpm_reset_ats =
|
||||
@@ -51,6 +54,7 @@ pub(super) async fn read_candidate_runtime_selection_snapshot(
|
||||
provider_key_rpm_states,
|
||||
provider_quota_blocks_requests,
|
||||
key_account_quota_exhausted,
|
||||
key_oauth_invalid,
|
||||
provider_key_rpm_reset_ats,
|
||||
})
|
||||
}
|
||||
@@ -104,6 +108,11 @@ pub(super) fn is_candidate_selectable(
|
||||
.get(candidate.key_id.as_str())
|
||||
.copied()
|
||||
.unwrap_or(false),
|
||||
oauth_invalid: snapshot
|
||||
.key_oauth_invalid
|
||||
.get(candidate.key_id.as_str())
|
||||
.copied()
|
||||
.unwrap_or(false),
|
||||
rpm_reset_at: snapshot
|
||||
.provider_key_rpm_reset_ats
|
||||
.get(candidate.key_id.as_str())
|
||||
@@ -142,6 +151,11 @@ pub(super) fn current_candidate_runtime_skip_reason(
|
||||
.get(candidate.key_id.as_str())
|
||||
.copied()
|
||||
.unwrap_or(false),
|
||||
oauth_invalid: snapshot
|
||||
.key_oauth_invalid
|
||||
.get(candidate.key_id.as_str())
|
||||
.copied()
|
||||
.unwrap_or(false),
|
||||
rpm_reset_at,
|
||||
})
|
||||
}
|
||||
@@ -270,6 +284,40 @@ fn read_key_account_quota_exhaustion_map(
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn read_key_oauth_invalid_map(
|
||||
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
|
||||
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
|
||||
now_unix_secs: u64,
|
||||
) -> BTreeMap<String, bool> {
|
||||
candidates
|
||||
.iter()
|
||||
.map(|candidate| {
|
||||
let oauth_invalid = provider_key_rpm_states
|
||||
.get(candidate.key_id.as_str())
|
||||
.is_some_and(|key| key_requires_oauth_reauth(key, now_unix_secs));
|
||||
(candidate.key_id.clone(), oauth_invalid)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn key_requires_oauth_reauth(key: &StoredProviderCatalogKey, now_unix_secs: u64) -> bool {
|
||||
if !key.auth_type.trim().eq_ignore_ascii_case("oauth") {
|
||||
return false;
|
||||
}
|
||||
|
||||
let invalid_reason = key
|
||||
.oauth_invalid_reason
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if !invalid_reason.is_empty() {
|
||||
return !invalid_reason.starts_with("[REQUEST_FAILED]");
|
||||
}
|
||||
|
||||
key.expires_at_unix_secs
|
||||
.is_some_and(|value| value > 0 && value <= now_unix_secs)
|
||||
}
|
||||
|
||||
fn read_provider_key_rpm_reset_at_map(
|
||||
state: &(impl SchedulerRuntimeState + ?Sized),
|
||||
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
|
||||
|
||||
@@ -1241,6 +1241,143 @@ async fn skips_codex_candidate_when_account_quota_is_exhausted_and_pool_flag_ena
|
||||
assert_eq!(skipped[0].skip_reason, "account_quota_exhausted");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_oauth_invalid_candidate_before_local_auth_resolution() {
|
||||
let mut first = sample_row();
|
||||
first.provider_id = "provider-codex".to_string();
|
||||
first.provider_name = "codex".to_string();
|
||||
first.provider_type = "codex".to_string();
|
||||
first.endpoint_id = "endpoint-codex".to_string();
|
||||
first.endpoint_api_format = "openai:cli".to_string();
|
||||
first.key_id = "key-codex".to_string();
|
||||
first.key_name = "codex-invalid".to_string();
|
||||
first.key_auth_type = "oauth".to_string();
|
||||
first.key_api_formats = Some(vec!["openai:cli".to_string()]);
|
||||
first.key_global_priority_by_format = Some(serde_json::json!({"openai:cli": 1}));
|
||||
|
||||
let mut second = sample_row();
|
||||
second.provider_id = "provider-openai".to_string();
|
||||
second.provider_name = "openai".to_string();
|
||||
second.endpoint_id = "endpoint-openai".to_string();
|
||||
second.endpoint_api_format = "openai:cli".to_string();
|
||||
second.key_id = "key-openai".to_string();
|
||||
second.key_name = "fallback".to_string();
|
||||
second.key_api_formats = Some(vec!["openai:cli".to_string()]);
|
||||
second.key_global_priority_by_format = Some(serde_json::json!({"openai:cli": 2}));
|
||||
|
||||
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
first, second,
|
||||
]));
|
||||
let mut codex_provider = sample_provider("provider-codex", None);
|
||||
codex_provider.provider_type = "codex".to_string();
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![codex_provider, sample_provider("provider-openai", None)],
|
||||
Vec::new(),
|
||||
vec![
|
||||
{
|
||||
let mut key = sample_key("key-codex", "provider-codex", Some(10));
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
|
||||
key.oauth_invalid_reason = Some(
|
||||
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已被使用并轮换,请重新登录授权"
|
||||
.to_string(),
|
||||
);
|
||||
key
|
||||
},
|
||||
sample_key("key-openai", "provider-openai", Some(10)),
|
||||
],
|
||||
));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
|
||||
candidates,
|
||||
provider_catalog,
|
||||
quotas,
|
||||
request_candidates,
|
||||
),
|
||||
);
|
||||
|
||||
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
|
||||
state.data.as_ref(),
|
||||
&state,
|
||||
"openai:cli",
|
||||
"gpt-4.1",
|
||||
false,
|
||||
None,
|
||||
1_710_000_100,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
assert_eq!(selected.len(), 1);
|
||||
assert_eq!(selected[0].provider_id, "provider-openai");
|
||||
assert_eq!(skipped.len(), 1);
|
||||
assert_eq!(skipped[0].candidate.provider_id, "provider-codex");
|
||||
assert_eq!(skipped[0].skip_reason, "oauth_invalid");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn keeps_request_failed_oauth_candidate_selectable() {
|
||||
let mut row = sample_row();
|
||||
row.provider_id = "provider-codex".to_string();
|
||||
row.provider_name = "codex".to_string();
|
||||
row.provider_type = "codex".to_string();
|
||||
row.endpoint_id = "endpoint-codex".to_string();
|
||||
row.endpoint_api_format = "openai:cli".to_string();
|
||||
row.key_id = "key-codex".to_string();
|
||||
row.key_name = "codex-check-failed".to_string();
|
||||
row.key_auth_type = "oauth".to_string();
|
||||
row.key_api_formats = Some(vec!["openai:cli".to_string()]);
|
||||
|
||||
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
row,
|
||||
]));
|
||||
let mut provider = sample_provider("provider-codex", None);
|
||||
provider.provider_type = "codex".to_string();
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
Vec::new(),
|
||||
vec![{
|
||||
let mut key = sample_key("key-codex", "provider-codex", Some(10));
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
|
||||
key.oauth_invalid_reason = Some("[REQUEST_FAILED] 账号状态检查失败".to_string());
|
||||
key
|
||||
}],
|
||||
));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
|
||||
candidates,
|
||||
provider_catalog,
|
||||
quotas,
|
||||
request_candidates,
|
||||
),
|
||||
);
|
||||
|
||||
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
|
||||
state.data.as_ref(),
|
||||
&state,
|
||||
"openai:cli",
|
||||
"gpt-4.1",
|
||||
false,
|
||||
None,
|
||||
1_710_000_100,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
assert_eq!(selected.len(), 1);
|
||||
assert_eq!(selected[0].provider_id, "provider-codex");
|
||||
assert!(skipped.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn keeps_codex_candidate_selectable_when_exhausted_account_flag_is_disabled() {
|
||||
let mut first = sample_row();
|
||||
|
||||
Reference in New Issue
Block a user