fix(kiro): 放行可刷新 OAuth 调度候选 (#340)

2026-10-08 10:27:46 +08:00 · 2026-04-26 23:49:55 +08:00
parent f3a9fd4c82
commit e8763b3bbd
5 changed files with 300 additions and 21 deletions
@@ -1378,6 +1378,242 @@ async fn keeps_request_failed_oauth_candidate_selectable() {
    assert!(skipped.is_empty());
 }

+#[tokio::test]
+async fn keeps_refreshable_kiro_candidate_selectable_with_runtime_oauth_invalid_marker() {
+    let mut row = sample_row();
+    row.provider_id = "provider-kiro".to_string();
+    row.provider_name = "kiro".to_string();
+    row.provider_type = "kiro".to_string();
+    row.endpoint_id = "endpoint-kiro".to_string();
+    row.key_id = "key-kiro".to_string();
+    row.key_name = "kiro-refreshable".to_string();
+    row.key_auth_type = "oauth".to_string();
+
+    let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
+        row,
+    ]));
+    let mut provider = sample_provider("provider-kiro", None);
+    provider.provider_type = "kiro".to_string();
+    let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+        vec![provider],
+        Vec::new(),
+        vec![{
+            let mut key = sample_key("key-kiro", "provider-kiro", Some(10));
+            key.auth_type = "oauth".to_string();
+            key.encrypted_auth_config = Some("encrypted-refreshable-session".to_string());
+            key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
+            key.oauth_invalid_reason = Some("Kiro Token 无效或已过期".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:chat",
+        "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-kiro");
+    assert!(skipped.is_empty());
+}
+
+#[tokio::test]
+async fn keeps_refreshable_kiro_candidate_selectable_when_oauth_token_expired() {
+    let mut row = sample_row();
+    row.provider_id = "provider-kiro".to_string();
+    row.provider_name = "kiro".to_string();
+    row.provider_type = "kiro".to_string();
+    row.endpoint_id = "endpoint-kiro".to_string();
+    row.key_id = "key-kiro".to_string();
+    row.key_name = "kiro-expired-refreshable".to_string();
+    row.key_auth_type = "oauth".to_string();
+
+    let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
+        row,
+    ]));
+    let mut provider = sample_provider("provider-kiro", None);
+    provider.provider_type = "kiro".to_string();
+    let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+        vec![provider],
+        Vec::new(),
+        vec![{
+            let mut key = sample_key("key-kiro", "provider-kiro", Some(10));
+            key.auth_type = "oauth".to_string();
+            key.encrypted_auth_config = Some("encrypted-refreshable-session".to_string());
+            key.expires_at_unix_secs = Some(1_710_000_000);
+            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:chat",
+        "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-kiro");
+    assert!(skipped.is_empty());
+}
+
+#[tokio::test]
+async fn skips_kiro_candidate_after_refresh_failure_requires_reauth() {
+    let mut row = sample_row();
+    row.provider_id = "provider-kiro".to_string();
+    row.provider_name = "kiro".to_string();
+    row.provider_type = "kiro".to_string();
+    row.endpoint_id = "endpoint-kiro".to_string();
+    row.key_id = "key-kiro".to_string();
+    row.key_name = "kiro-refresh-failed".to_string();
+    row.key_auth_type = "oauth".to_string();
+
+    let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
+        row,
+    ]));
+    let mut provider = sample_provider("provider-kiro", None);
+    provider.provider_type = "kiro".to_string();
+    let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+        vec![provider],
+        Vec::new(),
+        vec![{
+            let mut key = sample_key("key-kiro", "provider-kiro", Some(10));
+            key.auth_type = "oauth".to_string();
+            key.encrypted_auth_config = Some("encrypted-refreshable-session".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
+        }],
+    ));
+    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:chat",
+        "gpt-4.1",
+        false,
+        None,
+        1_710_000_100,
+    )
+    .await
+    .expect("selection should succeed");
+
+    assert!(selected.is_empty());
+    assert_eq!(skipped.len(), 1);
+    assert_eq!(skipped[0].candidate.provider_id, "provider-kiro");
+    assert_eq!(skipped[0].skip_reason, "oauth_invalid");
+}
+
+#[tokio::test]
+async fn skips_refreshable_kiro_candidate_when_oauth_marker_is_account_block() {
+    let mut row = sample_row();
+    row.provider_id = "provider-kiro".to_string();
+    row.provider_name = "kiro".to_string();
+    row.provider_type = "kiro".to_string();
+    row.endpoint_id = "endpoint-kiro".to_string();
+    row.key_id = "key-kiro".to_string();
+    row.key_name = "kiro-account-blocked".to_string();
+    row.key_auth_type = "oauth".to_string();
+
+    let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
+        row,
+    ]));
+    let mut provider = sample_provider("provider-kiro", None);
+    provider.provider_type = "kiro".to_string();
+    let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+        vec![provider],
+        Vec::new(),
+        vec![{
+            let mut key = sample_key("key-kiro", "provider-kiro", Some(10));
+            key.auth_type = "oauth".to_string();
+            key.encrypted_auth_config = Some("encrypted-refreshable-session".to_string());
+            key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
+            key.oauth_invalid_reason = Some("账户已封禁: account banned".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:chat",
+        "gpt-4.1",
+        false,
+        None,
+        1_710_000_100,
+    )
+    .await
+    .expect("selection should succeed");
+
+    assert!(selected.is_empty());
+    assert_eq!(skipped.len(), 1);
+    assert_eq!(skipped[0].candidate.provider_id, "provider-kiro");
+    assert_eq!(skipped[0].skip_reason, "oauth_invalid");
+}
+
 #[tokio::test]
 async fn keeps_codex_candidate_selectable_when_exhausted_account_flag_is_disabled() {
    let mut first = sample_row();