diff --git a/.downstream/native-domains.md b/.downstream/native-domains.md index bcc457f4b1eb..798b4f4395ff 100644 --- a/.downstream/native-domains.md +++ b/.downstream/native-domains.md @@ -78,11 +78,16 @@ invented from an empty request. Reads never persist identities or timestamps. The page loads full identities automatically, limits detail reads to three, rejects obsolete page/policy/mode responses, and retains unsaved mode drafts. General account edits preserve the latest stored mode when it is omitted. -Existing seeds and persisted TUI identities are retained, and new accounts still +Existing seeds and persisted device identities are retained, and new accounts still default to `device`. -Enabled simulation uses one fixed credential-source TUI device with independent -real conversations. Stored `session` / `full` values and API inputs remain +Enabled simulation uses one randomly seed-derived, fixed credential-source device +with independent real conversations. The final outbound identity uses official +`codex_cli_rs` normalization for both originator and User-Agent, including the +per-account device path. Stored device fields, OS/architecture/terminal, effective +version and conversation identities retain their existing selection. Persisted +TUI values remain stored; normalization does not regenerate the device. +Stored `session` / `full` values and API inputs remain compatible but have effective mode `device`; the UI offers only `off` and fixed-device/multiple-conversation mode. Their stored raw values are not migrated. Every enabled mode scopes session/thread/parent/window identities one-to-one by diff --git a/.downstream/upstream-base b/.downstream/upstream-base index e73e4a1d6b4c..28599c895818 100644 --- a/.downstream/upstream-base +++ b/.downstream/upstream-base @@ -1 +1 @@ -v0.2.17 +v0.2.18 diff --git a/.downstream/upstream-review-v0.2.18.md b/.downstream/upstream-review-v0.2.18.md new file mode 100644 index 000000000000..8917be875a00 --- /dev/null +++ b/.downstream/upstream-review-v0.2.18.md @@ -0,0 +1,51 @@ +# Official v0.2.18 integration + +- Official tag: `v0.2.18`, commit `ac1b8f1271be0d674fa6f074403c1410f30a3a81`. +- Downstream starting commit: `48da9d2dd3c5cfceaaf5f8477556e32b83fd2852`. +- The operator explicitly selected the resolutions below before implementation. + Machine-readable risk and final unresolved-conflict status belong to + `upstream-risk.json`; source admission and production delivery are recorded + separately in the operator workspace. + +## Initial text conflicts and selected resolutions + +| Conflict files | Resolution | +| --- | --- | +| `backend/internal/repository/scheduler_cache.go` | Use the official memo/revision algorithm, original key names and batch Pipeline. Remove downstream v2 isolation, legacy-key cleanup and extra batch transactions. Retain independent account metadata projections and the official delete-account transaction. No cache migration, Redis clearing or startup barrier is added. | +| `backend/internal/service/gateway_forward.go` | Ordinary Claude SSE handling uses the official 403/529 helper. Retire the downstream SSE classifier, prelude buffering and same-account thinking repair; retain official HTTP 400 thinking/tool repair and historical diagnostic readers. | +| `backend/internal/service/gateway_forward_as_chat_completions.go`, `backend/internal/service/gemini_chat_completions_compat_service.go` | Adopt official disconnect/usage accounting and stream readers. Retire the downstream streaming tool consistency guard and its dedicated error handling in the shared converter and all callers. Retain independent buffered argument reconstruction by actual call ID. | +| `backend/internal/service/openai_gateway_upstream_errors.go` | Remove the downstream special HTTP 400 upstream-error-wrapper classifier and all dependent calls/annotations. Use official error classification and capacity Retry-After waiting. | +| `backend/internal/service/openai_oauth_passthrough_test.go`, `backend/internal/service/openai_setup_token_compat_test.go` | Assertions use the final per-account CLI identity. Retain fixed randomly derived devices, independent conversations, turn state and strict transport; reuse the official normalization helper at the final outbound selector. | +| `frontend/package.json`, `frontend/pnpm-lock.yaml` | Use official v0.2.18 overrides, DOMPurify 3.4.17 and follow-redirects 1.16.1; remove redundant nanoid/Vue constraints. Preserve the independent xlsx 0.20.3 author-distribution dependency. | +| `frontend/src/components/account/AccountUsageCell.vue` | Use the official Gemini/Vertex usage region and service-account quota capability. Remove the former common quota bars and loading-with-old-values condition in that region. | +| `frontend/src/components/account/CreateAccountModal.vue`, `frontend/src/components/account/EditAccountModal.vue`, `frontend/src/components/admin/account/AccountActionMenu.vue` | Merge the imports used by official quota support and the independent model-sync/system-prompt features. Adopt official Vertex quota editing/reset behavior. | +| `frontend/src/views/admin/AccountsView.vue` | Platform/status return to official single-select localStorage persistence. Ignore old plural platform/status URL conditions; keep other filters and sensitive session state. Restore the official timer verbatim with `immediate: false`; it remains inactive, as explicitly selected after checking the library startup behavior. | + +These rows cover all 14 initial conflict files. Clean textual merges do not +override these decisions: the final per-account identity, shared tool converter, +filter state and dependent error-classification calls follow the same contract. + +## Official additions and retained boundaries + +- Preserve official `246_enable_gemini_image_generation_groups.sql` unchanged + alongside historical `246_plugin_extension_state.sql`. Existing Gemini and + Antigravity groups are enabled for image generation. Composite groups retain + their stored switch and reject Gemini image requests when it is off. +- Adopt the remaining conflict-free v0.2.18 changes, including stream billing, + Vertex quota accounting, scheduling performance, Cloudflare 1010 handling, + Gemini tool schemas, DeepSeek image payloads and image-storage download rules. + Go, Docker and CI use official 1.27.3. +- Preserve independent system prompts, DeepSeek default effort, model capacity + evidence and source binding, administrator visibility, account jobs/taxonomy, + login protection, strict Codex transport and Pelican workbench behavior. + Retired borrowing and recovery interfaces remain absent. Historical SQL names, + checksums and diagnostic records remain intact. +- Use the existing required PR checks and fixed-image release/deploy flow. Local + checks are limited to the changed risks; successful checks for the same code + tree are reused. No new fallback, retry, recovery switch or compatibility layer + is introduced. No real model calls or live payment tests are part of this + integration. + +Source integration does not establish release or deployment completion. The +operator production pointer and delivery evidence remain authoritative for the +running version, immutable image digest and actual verification results. diff --git a/.downstream/upstream-risk.json b/.downstream/upstream-risk.json index f90d9ea31ae6..178c262cc047 100644 --- a/.downstream/upstream-risk.json +++ b/.downstream/upstream-risk.json @@ -1,54 +1,56 @@ { - "critical_file_count": 45, + "critical_file_count": 47, "critical_files": [ + ".github/workflows/backend-ci.yml", + ".github/workflows/release.yml", + ".github/workflows/security-scan.yml", + "Dockerfile", + "backend/go.mod", + "backend/internal/repository/scheduler_cache.go", + "backend/internal/repository/scheduler_cache_snapshot_memo_unit_test.go", + "backend/internal/repository/simple_mode_default_groups.go", + "backend/internal/repository/simple_mode_default_groups_integration_test.go", + "backend/internal/repository/usage_billing_repo.go", + "backend/internal/repository/usage_billing_repo_integration_test.go", + "backend/internal/repository/usage_billing_repo_unit_test.go", "backend/internal/service/billing_service.go", - "backend/internal/service/gateway_anthropic_base_url_test.go", - "backend/internal/service/gateway_anthropic_passthrough.go", - "backend/internal/service/gateway_count_tokens.go", + "backend/internal/service/billing_service_test.go", + "backend/internal/service/gateway_forward.go", "backend/internal/service/gateway_forward_as_chat_completions.go", + "backend/internal/service/gateway_forward_as_chat_completions_ping_test.go", "backend/internal/service/gateway_forward_as_responses.go", - "backend/internal/service/gateway_hotpath_optimization_test.go", - "backend/internal/service/gateway_model_availability.go", - "backend/internal/service/gateway_service.go", - "backend/internal/service/gateway_upstream_request.go", - "backend/internal/service/gateway_upstream_request_test.go", - "backend/internal/service/gateway_websearch_emulation.go", - "backend/internal/service/gateway_websearch_emulation_test.go", - "backend/internal/service/openai_account_scheduler.go", - "backend/internal/service/openai_account_scheduler_test.go", - "backend/internal/service/openai_codex_model_metadata.go", - "backend/internal/service/openai_codex_model_metadata_test.go", - "backend/internal/service/openai_codex_models_service.go", - "backend/internal/service/openai_compact_fallback.go", - "backend/internal/service/openai_gateway_chat_completions.go", + "backend/internal/service/gateway_forward_compat_client_disconnect_test.go", + "backend/internal/service/gateway_forward_compat_stream.go", + "backend/internal/service/gateway_scheduling.go", + "backend/internal/service/gateway_usage_billing.go", + "backend/internal/service/gateway_usage_billing_account_quota_test.go", + "backend/internal/service/openai_capacity_shed_test.go", + "backend/internal/service/openai_codex_identity.go", + "backend/internal/service/openai_codex_identity_test.go", + "backend/internal/service/openai_codex_version_sync_service_test.go", + "backend/internal/service/openai_gateway_anthropic_native_pump.go", + "backend/internal/service/openai_gateway_anthropic_native_pump_test.go", "backend/internal/service/openai_gateway_chat_completions_anthropic_native.go", - "backend/internal/service/openai_gateway_forward.go", - "backend/internal/service/openai_gateway_grok_chat_bridge.go", - "backend/internal/service/openai_gateway_messages.go", - "backend/internal/service/openai_gateway_messages_chat_fallback.go", - "backend/internal/service/openai_gateway_model_availability.go", + "backend/internal/service/openai_gateway_deepseek_input_image_test.go", "backend/internal/service/openai_gateway_passthrough.go", "backend/internal/service/openai_gateway_request_body.go", - "backend/internal/service/openai_gateway_response_flush_test.go", - "backend/internal/service/openai_gateway_response_handling.go", - "backend/internal/service/openai_gateway_scheduling.go", - "backend/internal/service/openai_model_retirement.go", - "backend/internal/service/openai_model_retirement_test.go", - "backend/internal/service/openai_models_list.go", - "backend/internal/service/openai_response_model_rewrite_test.go", - "backend/internal/service/openai_responses_created_at.go", - "backend/internal/service/openai_responses_created_at_test.go", - "backend/internal/service/openai_responses_ingress_compat.go", - "backend/internal/service/openai_ws_forwarder_ingress.go", - "backend/internal/service/openai_ws_forwarder_support.go", - "backend/internal/service/openai_ws_http_bridge.go", - "backend/internal/service/openai_ws_v2_passthrough_adapter.go", - "backend/internal/service/pricing_service.go", + "backend/internal/service/openai_gateway_responses_anthropic_native.go", + "backend/internal/service/openai_gateway_service.go", + "backend/internal/service/openai_gateway_upstream_errors.go", + "backend/internal/service/openai_oauth_passthrough_test.go", + "backend/internal/service/openai_setup_token_compat_test.go", + "backend/internal/service/openai_ws_forwarder_success_test.go", + "backend/internal/service/openai_ws_http_bridge_test.go", "backend/internal/service/pricing_service_test.go", - "deploy/config.example.yaml" + "backend/internal/service/ratelimit_service.go", + "backend/internal/service/ratelimit_service_403_html_test.go", + "backend/migrations/246_enable_gemini_image_generation_groups.sql", + "backend/migrations/gemini_image_generation_groups_migration_test.go", + "frontend/package.json", + "frontend/pnpm-lock.yaml" ], - "downstream_commit": "4ebe80cf5c69408fbc6b43a86c730e44f2f03f66", - "downstream_file_count": 1383, + "downstream_commit": "48da9d2dd3c5cfceaaf5f8477556e32b83fd2852", + "downstream_file_count": 1422, "downstream_files": [ ".dockerignore", ".downstream/README.md", @@ -61,6 +63,7 @@ ".downstream/upstream-review-v0.2.14.md", ".downstream/upstream-review-v0.2.15.md", ".downstream/upstream-review-v0.2.16.md", + ".downstream/upstream-review-v0.2.17.md", ".downstream/upstream-risk.json", ".gitattributes", ".github/audit-exceptions.yml", @@ -188,11 +191,15 @@ "backend/internal/handler/admin/codex_error_diagnostics.go", "backend/internal/handler/admin/codex_fingerprint_handler.go", "backend/internal/handler/admin/codex_fingerprint_handler_test.go", + "backend/internal/handler/admin/deepseek_reasoning_handler.go", + "backend/internal/handler/admin/deepseek_reasoning_handler_test.go", "backend/internal/handler/admin/group_handler_simple_mode_test.go", "backend/internal/handler/admin/payment_handler_test.go", "backend/internal/handler/admin/pelican_background_handler.go", "backend/internal/handler/admin/pelican_background_handler_test.go", "backend/internal/handler/admin/pelican_fixture_test.go", + "backend/internal/handler/admin/pelican_gallery_handler.go", + "backend/internal/handler/admin/pelican_gallery_handler_test.go", "backend/internal/handler/admin/pelican_test_handler.go", "backend/internal/handler/admin/pelican_test_handler_test.go", "backend/internal/handler/admin/scheduled_test_handler.go", @@ -211,6 +218,7 @@ "backend/internal/handler/auth_handler.go", "backend/internal/handler/auth_login_protection.go", "backend/internal/handler/auth_login_protection_test.go", + "backend/internal/handler/composite_platform.go", "backend/internal/handler/dto/credentials_redact.go", "backend/internal/handler/dto/credentials_redact_test.go", "backend/internal/handler/dto/mappers.go", @@ -306,6 +314,7 @@ "backend/internal/pkg/claude/constants.go", "backend/internal/pkg/claude/effort_catalog.go", "backend/internal/pkg/claude/effort_catalog_test.go", + "backend/internal/pkg/codextransport/client_version_test.go", "backend/internal/pkg/codextransport/http1.go", "backend/internal/pkg/codextransport/http1_test.go", "backend/internal/pkg/codextransport/profile.go", @@ -433,6 +442,8 @@ "backend/internal/repository/pelican_background_repo_integration_test.go", "backend/internal/repository/pelican_deletion_repo.go", "backend/internal/repository/pelican_deletion_repo_integration_test.go", + "backend/internal/repository/pelican_gallery_repo_integration_test.go", + "backend/internal/repository/pelican_repeats_repo_integration_test.go", "backend/internal/repository/pelican_test_repo.go", "backend/internal/repository/pelican_test_repo_integration_test.go", "backend/internal/repository/pelican_test_repo_test.go", @@ -477,6 +488,7 @@ "backend/internal/server/routes/account_job_routes_test.go", "backend/internal/server/routes/admin.go", "backend/internal/server/routes/composite_platform_test.go", + "backend/internal/server/routes/deepseek_reasoning_routes_test.go", "backend/internal/server/routes/gateway.go", "backend/internal/server/routes/gateway_codex_models_test.go", "backend/internal/server/routes/gateway_group_fixture_test.go", @@ -568,6 +580,7 @@ "backend/internal/service/codex_fingerprint_account_view_test.go", "backend/internal/service/codex_fingerprint_policy_test.go", "backend/internal/service/codex_outbound_contract.go", + "backend/internal/service/codex_version_follow_test.go", "backend/internal/service/composite_platform_test.go", "backend/internal/service/concurrency_service.go", "backend/internal/service/content_moderation_runtime_cache_test.go", @@ -575,6 +588,8 @@ "backend/internal/service/conversation_standard_pricing_test.go", "backend/internal/service/deepseek_pricing.go", "backend/internal/service/deepseek_pricing_test.go", + "backend/internal/service/deepseek_reasoning.go", + "backend/internal/service/deepseek_reasoning_test.go", "backend/internal/service/domain_constants.go", "backend/internal/service/email_service_smtp_test.go", "backend/internal/service/gateway_anthropic_passthrough.go", @@ -625,6 +640,7 @@ "backend/internal/service/model_context_capacity.go", "backend/internal/service/model_context_capacity_catalog_test.go", "backend/internal/service/model_context_capacity_parser.go", + "backend/internal/service/model_context_capacity_parser_test.go", "backend/internal/service/model_context_capacity_reference_test.go", "backend/internal/service/model_context_capacity_test.go", "backend/internal/service/model_context_identity.go", @@ -767,6 +783,7 @@ "backend/internal/service/openai_model_not_supported.go", "backend/internal/service/openai_model_not_supported_test.go", "backend/internal/service/openai_model_reasoning_alias_test.go", + "backend/internal/service/openai_model_retirement_test.go", "backend/internal/service/openai_models_list.go", "backend/internal/service/openai_models_list_test.go", "backend/internal/service/openai_named_output_namespace_test.go", @@ -809,6 +826,7 @@ "backend/internal/service/openai_request_integrity_semantics.go", "backend/internal/service/openai_response_model_rewrite_test.go", "backend/internal/service/openai_responses_continuation.go", + "backend/internal/service/openai_responses_created_at_test.go", "backend/internal/service/openai_responses_ingress_compat.go", "backend/internal/service/openai_responses_ingress_compat_test.go", "backend/internal/service/openai_responses_input_compat.go", @@ -877,12 +895,21 @@ "backend/internal/service/payment_order_result_test.go", "backend/internal/service/pelican_account_generator.go", "backend/internal/service/pelican_account_generator_test.go", + "backend/internal/service/pelican_admission_test.go", "backend/internal/service/pelican_background.go", "backend/internal/service/pelican_background_test.go", + "backend/internal/service/pelican_evaluation.go", + "backend/internal/service/pelican_evaluation_test.go", "backend/internal/service/pelican_execution.go", "backend/internal/service/pelican_fixture_test.go", + "backend/internal/service/pelican_gallery.go", + "backend/internal/service/pelican_gallery_test.go", + "backend/internal/service/pelican_original_prompt_test.go", "backend/internal/service/pelican_preview.go", "backend/internal/service/pelican_preview_test.go", + "backend/internal/service/pelican_renderer.go", + "backend/internal/service/pelican_renderer_test.go", + "backend/internal/service/pelican_repeats_test.go", "backend/internal/service/pelican_result.go", "backend/internal/service/pelican_test_catalog.go", "backend/internal/service/pelican_test_catalog_test.go", @@ -1022,6 +1049,8 @@ "backend/migrations/277_drop_legacy_platform_check_constraints.sql", "backend/migrations/278_retire_codex_gateway_borrow_config.sql", "backend/migrations/279_pelican_test_deletion.sql", + "backend/migrations/280_pelican_evaluation.sql", + "backend/migrations/281_pelican_repeat_samples.sql", "backend/migrations/account_jobs_slim_migration_test.go", "backend/migrations/historical_cindy_null_total_migration_integration_test.go", "backend/migrations/minimax_platform_migration_test.go", @@ -1062,6 +1091,7 @@ "frontend/src/api/admin/accounts.ts", "frontend/src/api/admin/codexFingerprint.ts", "frontend/src/api/admin/codexRuntime.ts", + "frontend/src/api/admin/deepseekReasoning.ts", "frontend/src/api/admin/index.ts", "frontend/src/api/admin/loginProtection.ts", "frontend/src/api/admin/ops.ts", @@ -1164,12 +1194,16 @@ "frontend/src/components/admin/codex/CodexContinuationDiagnostics.vue", "frontend/src/components/admin/codex/CodexLayout.vue", "frontend/src/components/admin/codex/CodexRuntimeSettings.vue", + "frontend/src/components/admin/codex/PelicanAccountFilter.vue", + "frontend/src/components/admin/codex/PelicanEvaluationBadges.vue", "frontend/src/components/admin/codex/PelicanPreview.vue", "frontend/src/components/admin/codex/PelicanResultCard.vue", "frontend/src/components/admin/codex/PelicanTaskComposer.vue", + "frontend/src/components/admin/codex/PelicanWorkDetail.vue", "frontend/src/components/admin/codex/__tests__/CodexRuntime.spec.ts", "frontend/src/components/admin/codex/__tests__/PelicanPreview.spec.ts", "frontend/src/components/admin/codex/__tests__/PelicanTaskComposer.spec.ts", + "frontend/src/components/admin/codex/__tests__/PelicanWorkDetail.spec.ts", "frontend/src/components/admin/group/CodexManifestAccountsField.vue", "frontend/src/components/admin/group/GroupRPMOverridesModal.vue", "frontend/src/components/admin/group/GroupRateMultipliersModal.vue", @@ -1276,6 +1310,7 @@ "frontend/src/i18n/__tests__/accountConsoleLocales.spec.ts", "frontend/src/i18n/locales/en/admin/accountTasks.ts", "frontend/src/i18n/locales/en/admin/accounts.ts", + "frontend/src/i18n/locales/en/admin/deepseekReasoning.ts", "frontend/src/i18n/locales/en/admin/index.ts", "frontend/src/i18n/locales/en/admin/loginProtection.ts", "frontend/src/i18n/locales/en/admin/ops.ts", @@ -1288,6 +1323,7 @@ "frontend/src/i18n/locales/en/misc.ts", "frontend/src/i18n/locales/zh/admin/accountTasks.ts", "frontend/src/i18n/locales/zh/admin/accounts.ts", + "frontend/src/i18n/locales/zh/admin/deepseekReasoning.ts", "frontend/src/i18n/locales/zh/admin/index.ts", "frontend/src/i18n/locales/zh/admin/loginProtection.ts", "frontend/src/i18n/locales/zh/admin/ops.ts", @@ -1339,6 +1375,8 @@ "frontend/src/utils/keyGroupProviders.ts", "frontend/src/utils/modelContextCapacity.ts", "frontend/src/utils/openaiProtocolModeEdits.ts", + "frontend/src/utils/pelicanDisplay.ts", + "frontend/src/utils/pelicanDraft.ts", "frontend/src/utils/pelicanTaskSelection.ts", "frontend/src/utils/planType.ts", "frontend/src/utils/platformColors.ts", @@ -1359,6 +1397,7 @@ "frontend/src/views/admin/CodexPelicanComparisonView.vue", "frontend/src/views/admin/CodexRuntimeView.vue", "frontend/src/views/admin/DashboardView.vue", + "frontend/src/views/admin/DeepSeekReasoningView.vue", "frontend/src/views/admin/GroupsView.vue", "frontend/src/views/admin/LoginProtectionView.vue", "frontend/src/views/admin/PluginsView.vue", @@ -1382,6 +1421,7 @@ "frontend/src/views/admin/__tests__/AccountsView.usageWindowsHint.spec.ts", "frontend/src/views/admin/__tests__/CodexFingerprintView.spec.ts", "frontend/src/views/admin/__tests__/CodexPelicanComparisonView.spec.ts", + "frontend/src/views/admin/__tests__/DeepSeekReasoningView.spec.ts", "frontend/src/views/admin/__tests__/GroupsView.codexManifest.spec.ts", "frontend/src/views/admin/__tests__/GroupsView.duplicate.spec.ts", "frontend/src/views/admin/__tests__/LoginProtectionView.spec.ts", @@ -1432,273 +1472,207 @@ "frontend/src/workers/accountImport.worker.ts", "frontend/tailwind.config.js", "frontend/vite.config.ts", + "third_party/manxue-ai.LICENSE", "tools/check_pnpm_audit_exceptions.py" ], "initial_conflict_files": [ - "backend/internal/server/routes/composite_platform_test.go", - "backend/internal/service/account_test_service.go", - "backend/internal/service/antigravity_gateway_service.go", - "backend/internal/service/billing_service.go", - "backend/internal/service/gateway_model_availability.go", + "backend/internal/repository/scheduler_cache.go", + "backend/internal/service/gateway_forward.go", + "backend/internal/service/gateway_forward_as_chat_completions.go", "backend/internal/service/gemini_chat_completions_compat_service.go", - "backend/internal/service/gemini_messages_compat_service.go", - "backend/internal/service/openai_account_scheduler.go", - "backend/internal/service/openai_gateway_passthrough.go", - "backend/internal/service/openai_gateway_request_body.go", - "backend/internal/service/openai_gateway_response_handling.go", - "backend/internal/service/openai_gateway_scheduling.go", - "backend/internal/service/openai_ws_forwarder_ingress.go", - "backend/internal/service/openai_ws_forwarder_support.go", - "backend/internal/service/openai_ws_http_bridge.go", - "backend/internal/service/openai_ws_v2_passthrough_adapter.go", - "backend/internal/service/upstream_models.go", - "frontend/src/i18n/locales/en/admin/accounts.ts", - "frontend/src/i18n/locales/zh/admin/accounts.ts", - "frontend/src/views/admin/AccountsView.vue", - "frontend/src/views/user/__tests__/StripePaymentView.spec.ts" + "backend/internal/service/openai_gateway_upstream_errors.go", + "backend/internal/service/openai_oauth_passthrough_test.go", + "backend/internal/service/openai_setup_token_compat_test.go", + "frontend/package.json", + "frontend/pnpm-lock.yaml", + "frontend/src/components/account/AccountUsageCell.vue", + "frontend/src/components/account/CreateAccountModal.vue", + "frontend/src/components/account/EditAccountModal.vue", + "frontend/src/components/admin/account/AccountActionMenu.vue", + "frontend/src/views/admin/AccountsView.vue" ], "merge_conflicts": [], - "overlap_file_count": 80, + "overlap_file_count": 59, "overlap_files": [ + ".github/workflows/backend-ci.yml", + ".github/workflows/release.yml", + ".github/workflows/security-scan.yml", + "Dockerfile", "README.md", "README_CN.md", "README_JA.md", "backend/cmd/server/VERSION", - "backend/internal/config/config.go", - "backend/internal/config/config_test.go", - "backend/internal/handler/admin/account_codex_import.go", - "backend/internal/handler/admin/account_codex_import_test.go", - "backend/internal/handler/admin/account_handler.go", - "backend/internal/handler/admin/account_handler_available_models_test.go", + "backend/go.mod", + "backend/internal/handler/dto/mappers.go", "backend/internal/handler/gateway_handler.go", - "backend/internal/handler/gateway_models_test.go", - "backend/internal/handler/no_account_error.go", - "backend/internal/handler/no_account_error_test.go", - "backend/internal/handler/openai_gateway_handler.go", - "backend/internal/payment/provider/alipay_test.go", - "backend/internal/pkg/antigravity/client_test.go", - "backend/internal/pkg/antigravity/oauth_test.go", - "backend/internal/pkg/antigravity/request_transformer_test.go", - "backend/internal/pkg/apicompat/anthropic_to_responses.go", - "backend/internal/pkg/apicompat/chatcompletions_responses_test.go", - "backend/internal/pkg/apicompat/chatcompletions_to_responses.go", - "backend/internal/pkg/apicompat/types.go", - "backend/internal/server/routes/composite_platform_test.go", - "backend/internal/server/routes/gateway.go", - "backend/internal/server/routes/gateway_models_pinned_test.go", + "backend/internal/handler/gateway_handler_chat_completions.go", + "backend/internal/handler/gateway_handler_responses.go", + "backend/internal/repository/scheduler_cache.go", "backend/internal/service/account.go", - "backend/internal/service/account_test_service.go", + "backend/internal/service/admin_service_group_test.go", "backend/internal/service/antigravity_gateway_gemini.go", "backend/internal/service/antigravity_gateway_retry.go", - "backend/internal/service/antigravity_gateway_service.go", "backend/internal/service/billing_service.go", - "backend/internal/service/composite_platform_test.go", - "backend/internal/service/gateway_anthropic_passthrough.go", + "backend/internal/service/gateway_forward.go", "backend/internal/service/gateway_forward_as_chat_completions.go", "backend/internal/service/gateway_forward_as_responses.go", - "backend/internal/service/gateway_hotpath_optimization_test.go", - "backend/internal/service/gateway_model_availability.go", - "backend/internal/service/gateway_service.go", + "backend/internal/service/gateway_scheduling.go", + "backend/internal/service/gateway_usage_billing.go", "backend/internal/service/gemini_chat_completions_compat_service.go", "backend/internal/service/gemini_messages_compat_service.go", - "backend/internal/service/openai_account_scheduler.go", - "backend/internal/service/openai_account_scheduler_test.go", - "backend/internal/service/openai_codex_model_metadata.go", - "backend/internal/service/openai_codex_model_metadata_test.go", - "backend/internal/service/openai_codex_models_service.go", - "backend/internal/service/openai_gateway_chat_completions.go", + "backend/internal/service/gemini_messages_compat_service_test.go", + "backend/internal/service/model_rate_limit.go", + "backend/internal/service/openai_codex_identity.go", + "backend/internal/service/openai_codex_identity_test.go", + "backend/internal/service/openai_codex_version_sync_service_test.go", + "backend/internal/service/openai_gateway_anthropic_native_pump_test.go", "backend/internal/service/openai_gateway_chat_completions_anthropic_native.go", - "backend/internal/service/openai_gateway_forward.go", - "backend/internal/service/openai_gateway_grok_chat_bridge.go", - "backend/internal/service/openai_gateway_messages.go", - "backend/internal/service/openai_gateway_messages_chat_fallback.go", - "backend/internal/service/openai_gateway_model_availability.go", "backend/internal/service/openai_gateway_passthrough.go", "backend/internal/service/openai_gateway_request_body.go", - "backend/internal/service/openai_gateway_response_flush_test.go", - "backend/internal/service/openai_gateway_response_handling.go", - "backend/internal/service/openai_gateway_scheduling.go", - "backend/internal/service/openai_models_list.go", - "backend/internal/service/openai_response_model_rewrite_test.go", - "backend/internal/service/openai_responses_ingress_compat.go", - "backend/internal/service/openai_ws_forwarder_ingress.go", - "backend/internal/service/openai_ws_forwarder_support.go", - "backend/internal/service/openai_ws_http_bridge.go", - "backend/internal/service/openai_ws_v2_passthrough_adapter.go", - "backend/internal/service/pricing_service.go", - "backend/internal/service/upstream_models.go", + "backend/internal/service/openai_gateway_responses_anthropic_native.go", + "backend/internal/service/openai_gateway_service.go", + "backend/internal/service/openai_gateway_upstream_errors.go", + "backend/internal/service/openai_oauth_passthrough_test.go", + "backend/internal/service/openai_setup_token_compat_test.go", + "backend/internal/service/openai_ws_forwarder_success_test.go", + "backend/internal/service/openai_ws_http_bridge_test.go", + "backend/internal/service/ratelimit_service.go", + "backend/internal/service/setting_gateway_runtime.go", + "backend/internal/service/settings_view.go", "backend/resources/model-pricing/model_prices_and_context_window.json", - "deploy/config.example.yaml", - "frontend/src/components/account/AccountTestModal.vue", + "frontend/package.json", + "frontend/pnpm-lock.yaml", + "frontend/src/components/account/AccountCapacityCell.vue", + "frontend/src/components/account/AccountUsageCell.vue", + "frontend/src/components/account/CreateAccountModal.vue", "frontend/src/components/account/EditAccountModal.vue", - "frontend/src/components/admin/account/AccountTestModal.vue", - "frontend/src/composables/__tests__/useModelWhitelist.spec.ts", - "frontend/src/composables/useModelWhitelist.ts", - "frontend/src/i18n/locales/en/admin/accounts.ts", - "frontend/src/i18n/locales/zh/admin/accounts.ts", + "frontend/src/components/account/__tests__/AccountUsageCell.spec.ts", + "frontend/src/components/account/__tests__/EditAccountModal.spec.ts", + "frontend/src/components/admin/account/AccountActionMenu.vue", + "frontend/src/i18n/locales/en/admin/settings.ts", + "frontend/src/i18n/locales/zh/admin/settings.ts", "frontend/src/views/admin/AccountsView.vue", - "frontend/src/views/user/PaymentView.vue", - "frontend/src/views/user/StripePaymentView.vue", - "frontend/src/views/user/__tests__/StripePaymentView.spec.ts" + "frontend/src/views/admin/GroupsView.vue" ], "risk_class": "review_required", "schema": 1, - "upstream_base": "v0.2.16", - "upstream_file_count": 153, + "upstream_base": "v0.2.17", + "upstream_file_count": 114, "upstream_files": [ + ".github/workflows/backend-ci.yml", + ".github/workflows/release.yml", + ".github/workflows/security-scan.yml", + "Dockerfile", "README.md", "README_CN.md", "README_JA.md", - "assets/partners/logos/aimzoon.jpg", - "assets/partners/logos/modelflare.png", - "assets/partners/logos/proxy4free.png", + "assets/partners/logos/chenju_ai.png", + "assets/partners/logos/ooko.ai.png", "backend/cmd/server/VERSION", - "backend/internal/config/config.go", - "backend/internal/config/config_test.go", - "backend/internal/domain/constants.go", - "backend/internal/domain/constants_test.go", - "backend/internal/handler/admin/account_codex_import.go", - "backend/internal/handler/admin/account_codex_import_test.go", - "backend/internal/handler/admin/account_handler.go", - "backend/internal/handler/admin/account_handler_available_models_test.go", + "backend/go.mod", + "backend/internal/handler/dto/account_quota_service_account_test.go", + "backend/internal/handler/dto/mappers.go", + "backend/internal/handler/gateway_compat_stream_usage_test.go", "backend/internal/handler/gateway_handler.go", - "backend/internal/handler/gateway_models_test.go", + "backend/internal/handler/gateway_handler_chat_completions.go", + "backend/internal/handler/gateway_handler_responses.go", + "backend/internal/handler/gateway_handler_warmup_intercept_unit_test.go", + "backend/internal/handler/gemini_client_cancel_test.go", + "backend/internal/handler/gemini_image_generation_gate.go", + "backend/internal/handler/gemini_image_generation_gate_test.go", "backend/internal/handler/gemini_v1beta_handler.go", - "backend/internal/handler/gemini_v1beta_handler_test.go", - "backend/internal/handler/no_account_error.go", - "backend/internal/handler/no_account_error_test.go", - "backend/internal/handler/openai_gateway_handler.go", - "backend/internal/handler/payment_webhook_handler.go", - "backend/internal/handler/payment_webhook_handler_test.go", - "backend/internal/payment/provider/alipay.go", - "backend/internal/payment/provider/alipay_test.go", - "backend/internal/payment/provider/easypay.go", - "backend/internal/payment/provider/easypay_confirm_test.go", - "backend/internal/payment/provider/easypay_query_test.go", - "backend/internal/payment/types.go", - "backend/internal/pkg/antigravity/claude_types.go", - "backend/internal/pkg/antigravity/claude_types_test.go", - "backend/internal/pkg/antigravity/client.go", - "backend/internal/pkg/antigravity/client_test.go", - "backend/internal/pkg/antigravity/oauth.go", - "backend/internal/pkg/antigravity/oauth_test.go", - "backend/internal/pkg/antigravity/request_transformer.go", - "backend/internal/pkg/antigravity/request_transformer_test.go", - "backend/internal/pkg/apicompat/anthropic_to_responses.go", - "backend/internal/pkg/apicompat/chatcompletions_anthropic_bridge.go", - "backend/internal/pkg/apicompat/chatcompletions_anthropic_bridge_test.go", - "backend/internal/pkg/apicompat/chatcompletions_responses_test.go", - "backend/internal/pkg/apicompat/chatcompletions_to_responses.go", - "backend/internal/pkg/apicompat/types.go", - "backend/internal/pkg/gemini/models.go", - "backend/internal/pkg/geminicli/models.go", - "backend/internal/server/routes/composite_platform_test.go", - "backend/internal/server/routes/gateway.go", - "backend/internal/server/routes/gateway_models_pinned_test.go", + "backend/internal/handler/single_account_retry_probe_unit_test.go", + "backend/internal/pkg/antigravity/response_transformer.go", + "backend/internal/pkg/antigravity/stream_transformer.go", + "backend/internal/pkg/antigravity/usage_contract_test.go", + "backend/internal/pkg/openai/request.go", + "backend/internal/pkg/openai/request_codex_identity_normalization_test.go", + "backend/internal/repository/scheduler_cache.go", + "backend/internal/repository/scheduler_cache_snapshot_memo_unit_test.go", + "backend/internal/repository/simple_mode_default_groups.go", + "backend/internal/repository/simple_mode_default_groups_integration_test.go", + "backend/internal/repository/usage_billing_repo.go", + "backend/internal/repository/usage_billing_repo_integration_test.go", + "backend/internal/repository/usage_billing_repo_unit_test.go", "backend/internal/service/account.go", - "backend/internal/service/account_test_service.go", - "backend/internal/service/account_test_service_antigravity_test.go", - "backend/internal/service/account_websearch_test.go", + "backend/internal/service/account_model_mapping_signature_test.go", + "backend/internal/service/account_quota_schedulable_test.go", + "backend/internal/service/account_quota_service_account_test.go", + "backend/internal/service/account_shared_snapshot_test.go", + "backend/internal/service/admin_group.go", + "backend/internal/service/admin_service_group_test.go", + "backend/internal/service/antigravity_gateway_claude.go", "backend/internal/service/antigravity_gateway_compat.go", "backend/internal/service/antigravity_gateway_gemini.go", "backend/internal/service/antigravity_gateway_retry.go", - "backend/internal/service/antigravity_gateway_service.go", - "backend/internal/service/antigravity_gateway_service_test.go", - "backend/internal/service/antigravity_gateway_streaming.go", - "backend/internal/service/antigravity_gemini_thinking_variant.go", - "backend/internal/service/antigravity_gemini_thinking_variant_test.go", - "backend/internal/service/antigravity_rate_limit_test.go", - "backend/internal/service/apicompat_config.go", "backend/internal/service/billing_service.go", - "backend/internal/service/cn_provider_quota_service.go", - "backend/internal/service/cn_providers_test.go", - "backend/internal/service/composite_model_route.go", - "backend/internal/service/composite_platform_test.go", - "backend/internal/service/composite_route_resolver.go", - "backend/internal/service/composite_route_resolver_test.go", - "backend/internal/service/gateway_anthropic_base_url_test.go", - "backend/internal/service/gateway_anthropic_passthrough.go", - "backend/internal/service/gateway_count_tokens.go", + "backend/internal/service/billing_service_test.go", + "backend/internal/service/gateway_forward.go", "backend/internal/service/gateway_forward_as_chat_completions.go", + "backend/internal/service/gateway_forward_as_chat_completions_ping_test.go", "backend/internal/service/gateway_forward_as_responses.go", - "backend/internal/service/gateway_hotpath_optimization_test.go", - "backend/internal/service/gateway_model_availability.go", - "backend/internal/service/gateway_service.go", - "backend/internal/service/gateway_upstream_request.go", - "backend/internal/service/gateway_upstream_request_test.go", - "backend/internal/service/gateway_websearch_emulation.go", - "backend/internal/service/gateway_websearch_emulation_test.go", - "backend/internal/service/gemini_cache_storage_pricing_test.go", + "backend/internal/service/gateway_forward_compat_client_disconnect_test.go", + "backend/internal/service/gateway_forward_compat_stream.go", + "backend/internal/service/gateway_scheduling.go", + "backend/internal/service/gateway_usage_billing.go", + "backend/internal/service/gateway_usage_billing_account_quota_test.go", "backend/internal/service/gemini_chat_completions_compat_service.go", - "backend/internal/service/gemini_compat_thinking.go", - "backend/internal/service/gemini_compat_thinking_test.go", - "backend/internal/service/gemini_error_policy_single_check_test.go", - "backend/internal/service/gemini_error_policy_test.go", + "backend/internal/service/gemini_image_generation_gate.go", + "backend/internal/service/gemini_image_generation_gate_test.go", "backend/internal/service/gemini_messages_compat_service.go", - "backend/internal/service/gemini_native_reasoning_pricing_test.go", - "backend/internal/service/gemini_vertex_api_key_forward_test.go", - "backend/internal/service/gemini_vertex_token_error.go", - "backend/internal/service/gemini_vertex_token_error_test.go", - "backend/internal/service/openai_account_scheduler.go", - "backend/internal/service/openai_account_scheduler_test.go", - "backend/internal/service/openai_codex_model_metadata.go", - "backend/internal/service/openai_codex_model_metadata_test.go", - "backend/internal/service/openai_codex_models_service.go", - "backend/internal/service/openai_compact_fallback.go", - "backend/internal/service/openai_gateway_chat_completions.go", + "backend/internal/service/gemini_messages_compat_service_test.go", + "backend/internal/service/gemini_native_response_signal_test.go", + "backend/internal/service/gemini_stream_client_abort.go", + "backend/internal/service/gemini_stream_client_abort_test.go", + "backend/internal/service/image_storage.go", + "backend/internal/service/image_storage_test.go", + "backend/internal/service/model_rate_limit.go", + "backend/internal/service/model_rate_limit_snapshot_cow_test.go", + "backend/internal/service/openai_capacity_shed_test.go", + "backend/internal/service/openai_codex_identity.go", + "backend/internal/service/openai_codex_identity_test.go", + "backend/internal/service/openai_codex_version_sync_service_test.go", + "backend/internal/service/openai_gateway_anthropic_native_pump.go", + "backend/internal/service/openai_gateway_anthropic_native_pump_test.go", "backend/internal/service/openai_gateway_chat_completions_anthropic_native.go", - "backend/internal/service/openai_gateway_forward.go", - "backend/internal/service/openai_gateway_grok_chat_bridge.go", - "backend/internal/service/openai_gateway_messages.go", - "backend/internal/service/openai_gateway_messages_chat_fallback.go", - "backend/internal/service/openai_gateway_model_availability.go", + "backend/internal/service/openai_gateway_deepseek_input_image_test.go", "backend/internal/service/openai_gateway_passthrough.go", "backend/internal/service/openai_gateway_request_body.go", - "backend/internal/service/openai_gateway_response_flush_test.go", - "backend/internal/service/openai_gateway_response_handling.go", - "backend/internal/service/openai_gateway_scheduling.go", - "backend/internal/service/openai_model_retirement.go", - "backend/internal/service/openai_model_retirement_test.go", - "backend/internal/service/openai_models_list.go", - "backend/internal/service/openai_response_model_rewrite_test.go", - "backend/internal/service/openai_responses_created_at.go", - "backend/internal/service/openai_responses_created_at_test.go", - "backend/internal/service/openai_responses_ingress_compat.go", - "backend/internal/service/openai_ws_forwarder_ingress.go", - "backend/internal/service/openai_ws_forwarder_support.go", - "backend/internal/service/openai_ws_http_bridge.go", - "backend/internal/service/openai_ws_v2_passthrough_adapter.go", - "backend/internal/service/payment_fulfillment.go", - "backend/internal/service/payment_notification_confirm_test.go", - "backend/internal/service/payment_order.go", - "backend/internal/service/payment_order_type_test.go", - "backend/internal/service/payment_refund.go", - "backend/internal/service/payment_refund_pending_hold_test.go", - "backend/internal/service/payment_refund_test.go", - "backend/internal/service/payment_resume_service.go", - "backend/internal/service/payment_resume_service_test.go", - "backend/internal/service/payment_service.go", - "backend/internal/service/pricing_service.go", + "backend/internal/service/openai_gateway_responses_anthropic_native.go", + "backend/internal/service/openai_gateway_service.go", + "backend/internal/service/openai_gateway_upstream_errors.go", + "backend/internal/service/openai_oauth_passthrough_test.go", + "backend/internal/service/openai_setup_token_compat_test.go", + "backend/internal/service/openai_ws_forwarder_success_test.go", + "backend/internal/service/openai_ws_http_bridge_test.go", "backend/internal/service/pricing_service_test.go", - "backend/internal/service/upstream_models.go", - "backend/internal/service/upstream_models_antigravity_test.go", - "backend/internal/service/vertex_service_account.go", + "backend/internal/service/ratelimit_service.go", + "backend/internal/service/ratelimit_service_403_html_test.go", + "backend/internal/service/request_metadata.go", + "backend/internal/service/scheduler_cache.go", + "backend/internal/service/setting_gateway_runtime.go", + "backend/internal/service/settings_view.go", + "backend/internal/service/single_account_retry_probe_test.go", + "backend/internal/service/usage_billing.go", + "backend/migrations/246_enable_gemini_image_generation_groups.sql", + "backend/migrations/gemini_image_generation_groups_migration_test.go", "backend/resources/model-pricing/model_prices_and_context_window.json", - "deploy/config.example.yaml", - "docs/OPENAI_GPT55_RETIREMENT.md", - "frontend/src/components/account/AccountTestModal.vue", + "frontend/package.json", + "frontend/pnpm-lock.yaml", + "frontend/src/components/account/AccountCapacityCell.vue", + "frontend/src/components/account/AccountUsageCell.vue", + "frontend/src/components/account/CreateAccountModal.vue", "frontend/src/components/account/EditAccountModal.vue", - "frontend/src/components/admin/account/AccountTableFilters.vue", - "frontend/src/components/admin/account/AccountTestModal.vue", - "frontend/src/components/admin/account/__tests__/AccountTableFilters.spec.ts", - "frontend/src/composables/__tests__/useModelWhitelist.spec.ts", - "frontend/src/composables/useModelWhitelist.ts", - "frontend/src/i18n/locales/en/admin/accounts.ts", - "frontend/src/i18n/locales/zh/admin/accounts.ts", + "frontend/src/components/account/__tests__/AccountUsageCell.spec.ts", + "frontend/src/components/account/__tests__/EditAccountModal.spec.ts", + "frontend/src/components/admin/account/AccountActionMenu.vue", + "frontend/src/i18n/locales/en/admin/settings.ts", + "frontend/src/i18n/locales/zh/admin/settings.ts", + "frontend/src/utils/accountQuota.ts", "frontend/src/views/admin/AccountsView.vue", - "frontend/src/views/user/PaymentView.vue", - "frontend/src/views/user/StripePaymentView.vue", - "frontend/src/views/user/__tests__/StripePaymentView.spec.ts" + "frontend/src/views/admin/GroupsView.vue", + "frontend/src/views/admin/__tests__/groupsImagePricing.spec.ts", + "frontend/src/views/admin/groupsImagePricing.ts" ], - "upstream_tag": "v0.2.17" + "upstream_tag": "v0.2.18" } diff --git a/.github/scripts/test_agent_protocol.sh b/.github/scripts/test_agent_protocol.sh index 159b5c004036..ec81250350e0 100644 --- a/.github/scripts/test_agent_protocol.sh +++ b/.github/scripts/test_agent_protocol.sh @@ -16,4 +16,4 @@ if [[ ! -d backend/internal/web/dist/assets ]]; then fi go -C backend test -p=1 -tags unit,embed ./internal/web -run 'Test(Agent|FrontendServer_Middleware|ServeEmbeddedFrontend|OverrideFilesNeverReceiveImmutableCacheHeaders)' -count=1 go -C backend test -p=1 ./internal/pkg/apicompat -run 'Test(Agent|Stream_|.*Chat.*Responses|.*Responses.*Chat|ResponsesAgentMessageCapability)' -count=1 -go -C backend test -p=1 -tags unit ./internal/service ./internal/handler -run 'Test(Agent|ForwardResponses_.*(Chat|Reasoning)|ForwardAsAnthropic_ForceChatCompletions|ForwardAsRawChatCompletions|ForwardAsChatCompletionsForGrokStreaming|OpenAIRawStreamTerminalState|AnthropicChatCallersSurfaceProtocolFailure|GeminiChatCallerSurfacesProtocolFailure|AntigravityChatCallerSurfacesProtocolFailure|NativeAnthropicChatProtocolGuardPreservesDisconnectDrain|ResponsesAgentMessageRoutingPreservesChatEligibility|.*NoAccount|.*SelectionFailure)' -count=1 +go -C backend test -p=1 -tags unit ./internal/service ./internal/handler -run 'Test(Agent|ForwardResponses_.*(Chat|Reasoning)|ForwardAsAnthropic_ForceChatCompletions|ForwardAsRawChatCompletions|ForwardAsChatCompletionsForGrokStreaming|OpenAIRawStreamTerminalState|ResponsesAgentMessageRoutingPreservesChatEligibility|.*NoAccount|.*SelectionFailure)' -count=1 diff --git a/.github/workflows/backend-ci.yml b/.github/workflows/backend-ci.yml index 4f9dbdcd374f..9fc3976aac10 100644 --- a/.github/workflows/backend-ci.yml +++ b/.github/workflows/backend-ci.yml @@ -39,7 +39,7 @@ jobs: cache-dependency-path: backend/go.sum - name: Verify Go version run: | - go version | grep -q 'go1.27.2' + go version | grep -q 'go1.27.3' - name: Unit tests working-directory: backend run: make test-unit @@ -81,7 +81,7 @@ jobs: cache-dependency-path: backend/go.sum - name: Verify Go version run: | - go version | grep -q 'go1.27.2' + go version | grep -q 'go1.27.3' - name: golangci-lint uses: golangci/golangci-lint-action@v9 with: diff --git a/.github/workflows/downstream-verify.yml b/.github/workflows/downstream-verify.yml index aaa16c3790ba..41f7275954f4 100644 --- a/.github/workflows/downstream-verify.yml +++ b/.github/workflows/downstream-verify.yml @@ -47,9 +47,9 @@ jobs: - name: DeepSeek default reasoning contracts working-directory: backend run: go test -p=1 -tags unit ./internal/service ./internal/handler/admin ./internal/server/routes -run '^TestDeepSeekDefault' -count=1 - - name: Claude stream failure and signature recovery contracts + - name: Claude HTTP signature repair and historical Ops contracts working-directory: backend - run: go test -p=1 -tags unit ./internal/service ./internal/handler -run '^TestClaude(ClassifyStreamError|ThinkingSignatureError|Prelude_|Signature|StreamError|StreamFailureStatus|OpsErrorLogger)|^TestGatewayService_Forward_(PreOutputSSE|PostOutputSSE)' -count=1 + run: go test -p=1 -tags unit ./internal/service ./internal/handler -run '^TestClaude(ThinkingSignatureError|Signature|OpsErrorLogger)|^TestGatewayService_Forward_(PreOutputSSE|PostOutputSSE)' -count=1 - name: Independent Pelican contracts working-directory: backend run: | diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 7b9ab16f35ed..d07863715747 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -139,7 +139,7 @@ jobs: check-latest: false cache: false - name: Verify Go version - run: go version | grep -q 'go1.27.2' + run: go version | grep -q 'go1.27.3' - name: Restore target-specific Go cache id: go-cache uses: actions/cache/restore@v6 diff --git a/.github/workflows/security-scan.yml b/.github/workflows/security-scan.yml index 745ad5b6fea1..101ef2940030 100644 --- a/.github/workflows/security-scan.yml +++ b/.github/workflows/security-scan.yml @@ -26,7 +26,7 @@ jobs: cache-dependency-path: backend/go.sum - name: Verify Go version run: | - go version | grep -q 'go1.27.2' + go version | grep -q 'go1.27.3' - name: Run govulncheck working-directory: backend run: | diff --git a/DOWNSTREAM.md b/DOWNSTREAM.md index 5e4d643b89f4..b41ac6510689 100644 --- a/DOWNSTREAM.md +++ b/DOWNSTREAM.md @@ -5,7 +5,7 @@ This fork maintains the `codexrip` patch set over official Sub2API releases. ## Source baseline - Integrated official baseline: the release recorded in `.downstream/upstream-base`. -- Retained compatibility decisions: [upstream review](.downstream/upstream-review-v0.2.17.md). +- Retained compatibility decisions: [upstream review](.downstream/upstream-review-v0.2.18.md). - The operator workspace `docs/sub2api.md` owns the production version and image digest. A source merge or Release does not establish deployment completion. ## Settings and account operations @@ -25,6 +25,19 @@ Folder/tag filters, explicit field clearing and untouched-field preservation remain supported. A bulk-update item without a saved target account fails with `target_missing` rather than selecting a fresh set of accounts. +Platform and status filters follow the official single-select controls and +`account-filter-values` localStorage object (`platform`, `status`). The page has +one filter state for both controls and requests. Legacy `platforms` / `statuses` +URL conditions are ignored; the backend's existing plural-parameter API remains. +Other filters retain their current URL or session storage. The official +five-minute upstream billing-rate timer retains the official `immediate: false` +option and is not started by the page. + +Vertex service accounts use the official quota limits, progress and reset +behavior. The Gemini/Vertex usage region uses the official template rather than +the former downstream common quota bars and loading-with-old-values condition. +Upstream model synchronization and system-prompt binding remain independent. + The Codex runtime setting is one switch, request body compression for the streaming `/responses` turns of OpenAI OAuth accounts; its contract is in [settings, stored data and rollback](.downstream/native-domains.md). @@ -52,6 +65,13 @@ quota rows represent unlimited access; rows with three NULL limits are purged by the official migration. The three `238_*` migrations retain separate filenames and checksums. +Scheduler snapshots use the official decoded in-process memo, revision +invalidation, original key names and batch Pipeline writes. The downstream v2 +namespace, legacy-key cleanup and extra batch transactions are retired. Account +metadata still carries the fields required by the independent downstream +features. Startup and periodic snapshot rebuilds follow official behavior; +the update adds no cache migration or Redis clearing. + ## Model discovery and context capacity The account test picker reads raw upstream model IDs and applies the saved mapping @@ -98,23 +118,17 @@ regardless of the Codex remote/file selection. ## Claude streaming failures and signature recovery -Ordinary Claude Messages forwarding classifies SSE error events by their error -type, with a narrow thinking-signature fallback for third-party errors missing -the type. A stream error is not an HTTP permission rejection. In Ops attempt -JSON, `stream_error.upstream_status_code` retains the error's semantic status; -the optional `upstream_http_status_code` records the actual response status. -Existing records and HTTP error semantics remain compatible without a migration. - -When the existing signature rectifier permits it, ordinary Claude forwarding -can repair rejected thinking history on the same account before any message -event is committed. It buffers only `message_start`, up to 64 KiB, while sending -ping immediately. Committing the next message event or exceeding the bound -ends recovery eligibility. The first thinking repair has one attempt independent -of the initial request's ten-second backoff budget; the existing HTTP-only tool -fallback retains its original conditions. Automatic passthrough and other -provider paths do not gain this repair. Fingerprint-only recovery diagnostics -stay attached to the upstream attempt and do not replace its failure or expose -new thinking text or signature values. +Ordinary Claude Messages and Anthropic compatibility streams follow official +v0.2.18 SSE error handling: an `overloaded_error` before output is recorded as +529; other stream error events use 403. The downstream type-based SSE classifier, +pre-message buffering and same-account SSE thinking repair are retired. Official +HTTP 400 thinking-signature repair and its existing tool fallback remain. + +Historical Ops diagnostics, including the optional actual HTTP status and SSE +signature-recovery fields, remain readable. The update does not rewrite records +or add a new recovery path. The official compatibility reader collects observed +usage after client disconnect and reports upstream stream interruption according +to the selected protocol. ## OpenAI forwarding and model identity @@ -140,6 +154,14 @@ resolved forward target for scheduling. A Codex WebSocket window rollover remove the previous response ID without inferring an anchor from the old window; owner, identity and verified-replay checks remain in force. +Streaming Responses-to-Chat tool handling and its callers follow official +v0.2.18. The downstream tool-parameter/call-ID/final-snapshot consistency guard +and its dedicated error frames are retired. Independent buffered response +reconstruction still supplements completed arguments by the actual call ID. +The special downstream classification of HTTP 400 `upstream_error` wrappers is +also retired; error classification, account effects and capacity Retry-After +waiting follow official behavior. + ## Release and production deployment - `main` is admitted through the normal required PR checks. @@ -167,13 +189,14 @@ stripping, persistent rejection memory, SSE repair, missing-item-id repair and special account failovers are retired. The global recovery page and its `GET/PUT /api/v1/admin/reasoning-recovery` interface are removed. Stored settings, historical SQL and diagnostic records remain intact; ordinary Claude thinking -signature recovery is independent and remains available. +signature recovery for HTTP 400 remains independent and available. Cline and Command Code share the existing system-prompt insertion points. Connection-test reasoning options require declared model support and a wire path that retains the chosen effort; unsupported and `ultra` options remain -unavailable. The inline account filters retain multi-selection, tags, proxy, -plan and count controls in both console-theme modes. +unavailable. The inline account filters retain tags, proxy, plan and count +controls in both console-theme modes; platform/status selection is described +under account operations above. Official migration 242 removes platform CHECK constraints. Immutable downstream 251, 260 and 265 restore the quota constraint during a fresh installation, so @@ -186,17 +209,29 @@ Scheduled discovery may identify newer official releases. It cannot deploy production. Each integration records its selected official commit, conflicts and retained downstream contracts before a new immutable release is published. -`.downstream/upstream-risk.json` records the v0.2.16-to-v0.2.17 integration against -downstream commit `4ebe80cf5c69408fbc6b43a86c730e44f2f03f66`. The risk gate -recomputes those exact inputs. `initial_conflict_files` records 21 resolved -conflicts; `merge_conflicts: []` records that none remain. The candidate requires -a completed review and the `upstream-reviewed` label before promotion. +`.downstream/upstream-risk.json` records the v0.2.17-to-v0.2.18 integration against +downstream commit `48da9d2dd3c5cfceaaf5f8477556e32b83fd2852`. The risk gate +recomputes those exact inputs. The 14 initial conflict files and operator-selected +resolutions are recorded in the [v0.2.18 review](.downstream/upstream-review-v0.2.18.md). +The candidate requires a completed review, no unresolved conflicts and the +`upstream-reviewed` label before promotion. + +Vertex quota limits, disconnected-stream billing, shared snapshot reuse, +Cloudflare 1010 handling, Gemini tool-schema fixes, DeepSeek image normalization +and public-only image-storage downloads follow v0.2.18. Official migration +`246_enable_gemini_image_generation_groups.sql` enables image generation on +existing Gemini/Antigravity groups. Composite groups retain their stored switch; +Gemini image requests are rejected when that switch is off. The migration keeps +its original filename and checksum alongside `246_plugin_extension_state.sql`. +Go follows official 1.27.3. Frontend overrides follow official v0.2.18, including +DOMPurify 3.4.17 and follow-redirects 1.16.1; redundant nanoid/Vue constraints are +removed, while the independent xlsx 0.20.3 author-distribution dependency remains. Gemini 4 Argon, thinking summaries, Vertex error handling, payment fixes and workspace-aware Codex imports follow v0.2.17. GPT-5.5 retirement applies to ChatGPT subscription upstream targets; API keys and explicit replacement mappings -remain supported. Account filters are inline in both console-theme modes, retaining -downstream multi-selection and taxonomy. `gateway.min_output_tokens` defaults to +remain supported. Account filters remain inline in both console-theme modes. +`gateway.min_output_tokens` defaults to 128, with a minimum of 16. Antigravity capacity observations use the shared source-bound evidence resolver, including provider-specific `maxTokens`. diff --git a/Dockerfile b/Dockerfile index 0d7d813cc0a1..a49c162183cc 100644 --- a/Dockerfile +++ b/Dockerfile @@ -8,7 +8,7 @@ # ============================================================================= ARG NODE_IMAGE=node:24-alpine -ARG GOLANG_IMAGE=golang:1.27.2-alpine +ARG GOLANG_IMAGE=golang:1.27.3-alpine ARG ALPINE_IMAGE=alpine:3.21 ARG POSTGRES_IMAGE=postgres:18-alpine ARG GOPROXY=https://goproxy.cn,direct diff --git a/README.md b/README.md index fb54bfaf1780..a2573aab48eb 100644 --- a/README.md +++ b/README.md @@ -151,6 +151,16 @@ Please read the following carefully before using this project: Modelflare: Full strength, stable, nothing watered down. Global SOTA models at a lower cost. Operated by US-based Havenbyte LLC, Modelflare supports GPT, Claude, Gemini, Grok, and other leading models through common compatible protocols, plus 40+ image and video models. Cache compensation is a highlight: selected OpenAI groups (price-first, stable, and premium) guarantee daily cache hit rates of 65%, 75%, or 85%, and if eligible requests fall short, Modelflare refunds the difference, visible on the usage dashboard after settlement. Special offers: a single top-up of US$20 unlocks the GPT discount group at a 0.015x price multiplier, and a single top-up of US$50 unlocks the Claude discount group at a 0.03x multiplier. Register via this exclusive link to get started. + +Chenju +Thanks to Chenju for sponsoring this project! This is probably the only ad slot in this section that doesn't ask you to sign up for anything. Good open-source projects deserve to be paid for, so Chenju is chipping in a gift and raising a glass to the author. Find Chenju on X to chat about API relays, upstream providers, or just shoot the breeze. + + + +ooko.ai +🔑 Buying AI accounts at top prices | Gemini · Claude · GPT | Main and alt accounts all welcome, any channel wanted | Pay-per-usage supported | Seeking first-hand resources from major providers, I'll bring the KEY.
WeChat: Liuyou-zawd | Telegram: AACC8888888 + + ## Overview diff --git a/README_CN.md b/README_CN.md index bf72b9e5ac07..45e5a1c47145 100644 --- a/README_CN.md +++ b/README_CN.md @@ -154,6 +154,16 @@ Modelflare：满血稳定不掺水，低成本享受全球 SOTA 模型。站点由美国 LLC 主体 Havenbyte LLC 运营，支持 GPT、Claude、Gemini、Grok 等主流模型和常用兼容协议，还支持 40 多种图片和视频模型。缓存补偿是一大亮点：部分 OpenAI 价格优先、稳定和高级分组提供每日 65%、75%、85% 的缓存命中率保障，符合条件的请求没达到目标的，平台补偿差额，结算后可在用量看板查看。特惠活动：单笔充值 US$20 即可解锁 0.015 倍价格倍率的 GPT 特惠分组，单笔充值 US$50 即可解锁 0.03 倍价格倍率的 Claude 特惠分组。通过专属链接注册即可参与 + +Chenju +感谢 陈局 赞助了本项目！这大概是本栏唯一一个不让你注册的广告位。好的开源项目值得被付费，陈局来随个份子，敬作者一杯。X 上找 陈局，聊中转、聊上游、吹吹水。 + + + +ooko.ai +🔑 高价收 AI 账号｜Gemini · Claude · GPT｜大小号通收，各种渠道都要｜支持按量跑付｜寻一手大厂资源，我上 KEY 。
微信：Liuyou-zawd｜TG：AACC8888888 + + ## 项目概述 diff --git a/README_JA.md b/README_JA.md index 38593daf3829..0ef6c110d607 100644 --- a/README_JA.md +++ b/README_JA.md @@ -153,6 +153,16 @@ Modelflare：フルスペック、安定、水増しなし。世界の SOTA モデルを、より低いコストで。米国拠点の Havenbyte LLC が運営する Modelflare は、一般的な互換プロトコルを通じて GPT、Claude、Gemini、Grok などの主要モデルに対応し、さらに 40 以上の画像・動画モデルも利用できます。注目はキャッシュ補償です。対象の OpenAI グループ（価格優先・安定・プレミアム）は、1 日のキャッシュヒット率をそれぞれ 65%、75%、85% で保証します。対象リクエストがこの水準に達しない場合、Modelflare が差額を返金し、精算後に利用状況ダッシュボードで確認できます。特別オファー：US$20 を 1 回チャージすると、価格倍率 0.015 倍の GPT 割引グループが解放されます。US$50 を 1 回チャージすると、価格倍率 0.03 倍の Claude 割引グループが解放されます。こちらの専用リンクから登録して始めましょう。 + +Chenju +陈局（Chenju）様、本プロジェクトへのご支援ありがとうございます！これはおそらく、この欄で唯一「登録しなくていい」広告枠です。良いオープンソースには対価が払われるべき——というわけで、陈局がご祝儀を包んで作者に一杯おごります。X で 陈局 を見つけて、中継サービスや上流の話、雑談でもどうぞ。 + + + +ooko.ai +🔑 AI アカウント高価買取｜Gemini · Claude · GPT｜メイン・サブアカウントすべて買取、チャネル問わず歓迎｜従量制での支払いに対応｜大手の一次リソースを求めています、KEY は私が用意します。
WeChat：Liuyou-zawd｜Telegram：AACC8888888 + + ## 概要 diff --git a/assets/partners/logos/chenju_ai.png b/assets/partners/logos/chenju_ai.png new file mode 100644 index 000000000000..2e4affee5d60 Binary files /dev/null and b/assets/partners/logos/chenju_ai.png differ diff --git a/assets/partners/logos/ooko.ai.png b/assets/partners/logos/ooko.ai.png new file mode 100644 index 000000000000..888458340352 Binary files /dev/null and b/assets/partners/logos/ooko.ai.png differ diff --git a/backend/go.mod b/backend/go.mod index aa028c1bee7f..fa28e596a150 100644 --- a/backend/go.mod +++ b/backend/go.mod @@ -1,6 +1,6 @@ module github.com/Wei-Shaw/sub2api -go 1.27.2 +go 1.27.3 require ( entgo.io/ent v0.14.5 diff --git a/backend/internal/handler/dto/account_quota_service_account_test.go b/backend/internal/handler/dto/account_quota_service_account_test.go new file mode 100644 index 000000000000..1f6f1ae5104f --- /dev/null +++ b/backend/internal/handler/dto/account_quota_service_account_test.go @@ -0,0 +1,78 @@ +package dto + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +// service_account 与 apikey 同语义：设了总额度才输出 quota_limit / quota_used。 +func TestAccountFromServiceShallow_ServiceAccountQuota(t *testing.T) { + t.Run("no limit omits quota fields", func(t *testing.T) { + out := AccountFromServiceShallow(&service.Account{ + ID: 1, + Platform: service.PlatformGemini, + Type: service.AccountTypeServiceAccount, + Extra: map[string]any{"quota_used": 12.5}, + }) + require.Nil(t, out.QuotaUsed) + require.Nil(t, out.QuotaLimit) + require.Nil(t, out.QuotaDailyLimit) + require.Nil(t, out.QuotaWeeklyLimit) + + payload, err := json.Marshal(out) + require.NoError(t, err) + var decoded map[string]any + require.NoError(t, json.Unmarshal(payload, &decoded)) + require.NotContains(t, decoded, "quota_used") + require.NotContains(t, decoded, "quota_limit") + }) + + t.Run("with limits outputs limits and usage", func(t *testing.T) { + out := AccountFromServiceShallow(&service.Account{ + ID: 3, + Platform: service.PlatformGemini, + Type: service.AccountTypeServiceAccount, + Extra: map[string]any{ + "quota_used": 30.0, + "quota_limit": 100.0, + "quota_daily_limit": 10.0, + }, + }) + require.NotNil(t, out.QuotaLimit) + require.InDelta(t, 100.0, *out.QuotaLimit, 1e-9) + require.NotNil(t, out.QuotaUsed) + require.InDelta(t, 30.0, *out.QuotaUsed, 1e-9) + require.NotNil(t, out.QuotaDailyLimit) + require.NotNil(t, out.QuotaDailyUsed) + }) + + t.Run("lite list item keeps quota fields", func(t *testing.T) { + item := AccountListItemFromAccount(AccountFromServiceShallow(&service.Account{ + ID: 4, + Platform: service.PlatformGemini, + Type: service.AccountTypeServiceAccount, + Extra: map[string]any{"quota_used": 7.25, "quota_limit": 290.0}, + })) + payload, err := json.Marshal(item) + require.NoError(t, err) + var decoded map[string]any + require.NoError(t, json.Unmarshal(payload, &decoded)) + require.InDelta(t, 7.25, decoded["quota_used"], 1e-9) + require.InDelta(t, 290.0, decoded["quota_limit"], 1e-9) + }) +} + +// apikey 维持原语义：未设总额度时不输出 quota_used。 +func TestAccountFromServiceShallow_APIKeyWithoutLimitOmitsQuotaUsed(t *testing.T) { + out := AccountFromServiceShallow(&service.Account{ + ID: 5, + Type: service.AccountTypeAPIKey, + Extra: map[string]any{"quota_used": 5.0}, + }) + require.Nil(t, out.QuotaUsed) + require.Nil(t, out.QuotaLimit) +} diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index 435e29d751f1..87709ee40cc5 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -344,8 +344,8 @@ func AccountFromServiceShallow(a *service.Account) *Account { } } - // 提取账号配额限制（apikey / bedrock 类型有效） - if a.IsAPIKeyOrBedrock() { + // 提取账号配额限制（apikey / bedrock / service_account 类型有效） + if a.SupportsAccountQuota() { if limit := a.GetQuotaLimit(); limit > 0 { out.QuotaLimit = &limit used := a.GetQuotaUsed() diff --git a/backend/internal/handler/gateway_claude_stream_error_test.go b/backend/internal/handler/gateway_claude_stream_error_test.go index 4d37d53d36a0..47442bc73e9f 100644 --- a/backend/internal/handler/gateway_claude_stream_error_test.go +++ b/backend/internal/handler/gateway_claude_stream_error_test.go @@ -5,7 +5,6 @@ import ( "fmt" "net/http" "net/http/httptest" - "strings" "testing" "github.com/Wei-Shaw/sub2api/internal/model" @@ -15,47 +14,6 @@ import ( "github.com/tidwall/gjson" ) -func TestClaudeStreamErrorTerminal(t *testing.T) { - gin.SetMode(gin.TestMode) - for _, test := range []struct { - errorType string - message string - status int - }{ - {"invalid_request_error", "Invalid thinking signature", 400}, - {"overloaded_error", "Overloaded", 503}, - {"service_error", "Service temporarily unavailable", 502}, - } { - for _, started := range []bool{false, true} { - t.Run(fmt.Sprintf("%s/stream=%v", test.errorType, started), func(t *testing.T) { - recorder := httptest.NewRecorder() - c, _ := gin.CreateTestContext(recorder) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) - if started { - _, err := c.Writer.WriteString(partialMessageStartSSE) - require.NoError(t, err) - } - body := []byte(fmt.Sprintf(`{"type":"error","error":{"type":%q,"message":%q}}`, test.errorType, test.message)) - failoverErr := service.ClassifyClaudeStreamError(body).FailoverError(body) - (&GatewayHandler{}).handleFailoverExhausted(c, failoverErr, service.PlatformAnthropic, started) - if started { - require.Equal(t, 200, recorder.Code) - require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: message_start")) - require.Equal(t, 1, strings.Count(recorder.Body.String(), `"type":"error"`)) - parsed := parseOpsErrorResponse(recorder.Body.Bytes()) - require.Equal(t, test.errorType, parsed.ErrorType) - require.Equal(t, test.message, parsed.Message) - } else { - require.Equal(t, test.status, recorder.Code) - require.Equal(t, test.errorType, gjson.Get(recorder.Body.String(), "error.type").String()) - require.Equal(t, test.message, gjson.Get(recorder.Body.String(), "error.message").String()) - } - require.NotContains(t, recorder.Body.String(), "Upstream access forbidden") - }) - } - } -} - type claudeStreamPassthroughRepo struct { service.ErrorPassthroughRepository rules []*model.ErrorPassthroughRule @@ -77,7 +35,7 @@ func TestClaudeStreamErrorTerminal_PassthroughRulePrecedence(t *testing.T) { c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) body := []byte(`{"error":{"type":"overloaded_error","message":"Overloaded"}}`) - handler.handleFailoverExhausted(c, service.ClassifyClaudeStreamError(body).FailoverError(body), service.PlatformAnthropic, false) + handler.handleFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: 529, ResponseBody: body}, service.PlatformAnthropic, false) require.Equal(t, responseCode, recorder.Code) require.Equal(t, message, gjson.Get(recorder.Body.String(), "error.message").String()) require.Equal(t, true, c.GetBool(service.OpsSkipPassthroughKey)) @@ -142,7 +100,7 @@ func TestClaudeOpsErrorLogger_StreamErrorSemanticStatus(t *testing.T) { c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{{Platform: service.PlatformAnthropic, Kind: "stream_error", UpstreamStatusCode: test.status, UpstreamHTTPStatusCode: 200, Message: "provider failure"}}) _, err := c.Writer.WriteString(partialMessageStartSSE) require.NoError(t, err) - (&GatewayHandler{}).handleFailoverExhausted(c, service.ClassifyClaudeStreamError(body).FailoverError(body), service.PlatformAnthropic, true) + (&GatewayHandler{}).handleFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: 529, ResponseBody: body}, service.PlatformAnthropic, true) }) recorder := httptest.NewRecorder() router.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/v1/messages", nil)) diff --git a/backend/internal/handler/gateway_compat_stream_usage_test.go b/backend/internal/handler/gateway_compat_stream_usage_test.go new file mode 100644 index 000000000000..84f2e587d431 --- /dev/null +++ b/backend/internal/handler/gateway_compat_stream_usage_test.go @@ -0,0 +1,229 @@ +//go:build unit + +package handler + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type compatStreamUsageUpstream struct { + service.HTTPUpstream + response *http.Response + calls int +} + +func (u *compatStreamUsageUpstream) DoWithTLS(*http.Request, string, int64, int, *tlsfingerprint.Profile) (*http.Response, error) { + u.calls++ + return u.response, nil +} + +type compatStreamUsageBillingRepo struct { + service.UsageBillingRepository + mu sync.Mutex + commands []*service.UsageBillingCommand +} + +func (r *compatStreamUsageBillingRepo) Apply(ctx context.Context, cmd *service.UsageBillingCommand) (*service.UsageBillingApplyResult, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.commands = append(r.commands, cmd) + return &service.UsageBillingApplyResult{Applied: true}, nil +} + +const compatStreamUsagePartialStream = "event: message_start\n" + + `data: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4.5","usage":{"input_tokens":12,"output_tokens":1}}}` + "\n\n" + + "event: content_block_start\n" + + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n" + + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}` + "\n\n" + +const compatStreamUsageCompletion = "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0}` + "\n\n" + + "event: message_delta\n" + + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":7}}` + "\n\n" + + "event: message_stop\n" + + `data: {"type":"message_stop"}` + "\n\n" + +type compatStreamUsageRun struct { + rec *httptest.ResponseRecorder + upstream *compatStreamUsageUpstream + usageRepo *openAIWSUsageHandlerUsageLogRepoStub + billingRepo *compatStreamUsageBillingRepo +} + +// runCompatStreamUsageHandler 经 GatewayHandler 转发一次流式请求；上游写完 upstreamStream 后， +// closeUpstream 为 false 时保持连接不再发送数据。 +func runCompatStreamUsageHandler(t *testing.T, endpoint, upstreamStream string, closeUpstream bool) compatStreamUsageRun { + t.Helper() + cfg := &config.Config{Gateway: config.GatewayConfig{StreamDataIntervalTimeout: 1}} + groupID := int64(9300) + group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformAnthropic, Status: service.StatusActive, RateMultiplier: 1} + account := &service.Account{ + ID: 9301, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey, + Status: service.StatusActive, Schedulable: true, Concurrency: 1, + Credentials: map[string]any{"api_key": "test-key"}, + AccountGroups: []service.AccountGroup{{AccountID: 9301, GroupID: groupID}}, + } + snapshot := service.NewSchedulerSnapshotService(&fakeSchedulerCache{accounts: []*service.Account{account}}, nil, nil, nil, nil) + run := compatStreamUsageRun{ + rec: httptest.NewRecorder(), + usageRepo: &openAIWSUsageHandlerUsageLogRepoStub{created: make(chan *service.UsageLog, 2)}, + billingRepo: &compatStreamUsageBillingRepo{}, + } + billingCache := service.NewBillingCacheService(newHandlerInflightCache(100), nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billingCache.Stop) + + pr, pw := io.Pipe() + t.Cleanup(func() { _ = pr.Close(); _ = pw.Close() }) + go func() { + _, _ = io.WriteString(pw, upstreamStream) + if closeUpstream { + _ = pw.Close() + } + }() + run.upstream = &compatStreamUsageUpstream{response: &http.Response{ + StatusCode: http.StatusOK, Body: pr, + Header: http.Header{"X-Request-Id": []string{"compat-stream-usage"}}, + }} + gw := service.NewGatewayService( + nil, &fakeGroupRepo{group: group}, run.usageRepo, run.billingRepo, nil, nil, nil, nil, cfg, + snapshot, nil, service.NewBillingService(cfg, nil), nil, nil, nil, run.upstream, + &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, + ) + h := &GatewayHandler{ + gatewayService: gw, billingCacheService: billingCache, cfg: cfg, maxAccountSwitches: 1, + concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatClaude, 0), + } + + body := `{"model":"claude-sonnet-4.5","stream":true,"input":"hello"}` + if endpoint == "chat/completions" { + body = `{"model":"claude-sonnet-4.5","stream":true,"messages":[{"role":"user","content":"hello"}]}` + } + c, _ := gin.CreateTestContext(run.rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/"+endpoint, bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + apiKey := &service.APIKey{ID: 9302, UserID: 9303, GroupID: &groupID, Group: group, User: &service.User{ID: 9303, Balance: 100}} + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 9303}) + + if endpoint == "chat/completions" { + h.ChatCompletions(c) + } else { + h.Responses(c) + } + return run +} + +// compatStreamUsageMarkers 返回成功收尾与终止错误的标记。 +func compatStreamUsageMarkers(endpoint string) (success, failure string) { + if endpoint == "responses" { + return "event: response.completed", "event: response.failed" + } + return "[DONE]", `data: {"error":` +} + +func TestGatewayCompatibilityHandlers_StreamUsageRecordedOnceWithTerminalEvent(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, endpoint := range []string{"chat/completions", "responses"} { + for _, tc := range []struct { + name string + completion string + wantOutput int + }{ + {name: "complete", completion: compatStreamUsageCompletion, wantOutput: 7}, + {name: "idle_before_message_stop", wantOutput: 1}, + } { + t.Run(endpoint+"/"+tc.name, func(t *testing.T) { + successMarker, failureMarker := compatStreamUsageMarkers(endpoint) + wantMarker, denyMarker := successMarker, failureMarker + if tc.completion == "" { + wantMarker, denyMarker = failureMarker, successMarker + } + + run := runCompatStreamUsageHandler(t, endpoint, compatStreamUsagePartialStream+tc.completion, tc.completion != "") + + require.Equal(t, 1, run.upstream.calls) + out := run.rec.Body.String() + require.Contains(t, out, "partial") + require.Contains(t, out, wantMarker) + require.NotContains(t, out, denyMarker) + require.Equal(t, 1, strings.Count(out, wantMarker), "exactly one terminal event") + require.NotContains(t, out, `data: {"type":"error"`, "the handler must not append a fallback frame") + if tc.completion == "" { + require.Contains(t, out, "stream_timeout") + } + + run.billingRepo.mu.Lock() + defer run.billingRepo.mu.Unlock() + require.Len(t, run.billingRepo.commands, 1, "usage must be settled exactly once") + require.Equal(t, 12, run.billingRepo.commands[0].InputTokens) + require.Equal(t, tc.wantOutput, run.billingRepo.commands[0].OutputTokens) + require.Len(t, run.usageRepo.created, 1) + }) + } + } +} + +func TestGatewayCompatibilityHandlers_IdleBeforeAnyEventWritesStreamError(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, endpoint := range []string{"chat/completions", "responses"} { + t.Run(endpoint, func(t *testing.T) { + successMarker, failureMarker := compatStreamUsageMarkers(endpoint) + + run := runCompatStreamUsageHandler(t, endpoint, "", false) + + require.Equal(t, 1, run.upstream.calls) + require.Equal(t, http.StatusOK, run.rec.Code) + require.Equal(t, "text/event-stream", run.rec.Header().Get("Content-Type")) + out := run.rec.Body.String() + require.Equal(t, 1, strings.Count(out, failureMarker), "exactly one terminal event") + require.NotContains(t, out, `data: {"type":"error"`, "the handler must not append a fallback frame") + require.Contains(t, out, "stream_timeout") + require.NotContains(t, out, successMarker) + + run.billingRepo.mu.Lock() + defer run.billingRepo.mu.Unlock() + require.Empty(t, run.billingRepo.commands, "no usage was observed") + require.Empty(t, run.usageRepo.created) + }) + } +} + +func TestGatewayCompatibilityHandlers_UpstreamErrorEventAfterOutputIsNotBilled(t *testing.T) { + gin.SetMode(gin.TestMode) + upstreamError := "event: error\n" + `data: {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}` + "\n\n" + for _, endpoint := range []string{"chat/completions", "responses"} { + t.Run(endpoint, func(t *testing.T) { + successMarker, failureMarker := compatStreamUsageMarkers(endpoint) + + run := runCompatStreamUsageHandler(t, endpoint, compatStreamUsagePartialStream+upstreamError, true) + + require.Equal(t, 1, run.upstream.calls) + out := run.rec.Body.String() + require.Contains(t, out, "partial") + require.Equal(t, 1, strings.Count(out, failureMarker), "exactly one terminal event") + require.NotContains(t, out, `data: {"type":"error"`, "the handler must not append a fallback frame") + require.Contains(t, out, "Overloaded") + require.NotContains(t, out, successMarker) + + run.billingRepo.mu.Lock() + defer run.billingRepo.mu.Unlock() + require.Empty(t, run.billingRepo.commands, "an upstream error event is not billed while the client is online") + require.Empty(t, run.usageRepo.created) + }) + } +} diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index a26873efed60..09d674c98b96 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -228,6 +228,10 @@ func (h *GatewayHandler) Messages(c *gin.Context) { if rejectSystemOneOnlyPlatform(c, apiKey, h.errorResponse) { return } + if geminiImageGenerationRejected(c, apiKey, nil, reqModel, channelMapping.MappedModel) { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", service.ImageGenerationPermissionMessage()) + return + } if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolAnthropicMessages, reqModel, body); decision != nil && !decision.AllowNextStage { h.anthropicSecurityAuditError(c, decision) @@ -339,12 +343,9 @@ func (h *GatewayHandler) Messages(c *gin.Context) { if platform == service.PlatformGemini { fs := NewFailoverState(h.maxAccountSwitchesGemini, hasBoundSession) - // 单账号分组提前设置 SingleAccountRetry 标记，让 Service 层首次 503 就不设模型限流标记。 - // 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流，导致后续请求连续快速失败。 - if h.gatewayService.IsSingleAntigravityAccountGroup(c.Request.Context(), apiKey.GroupID) { - ctx := service.WithSingleAccountRetry(c.Request.Context(), true, h.metadataBridgeEnabled()) - c.Request = c.Request.WithContext(ctx) - } + // 单账号分组首次 503 就不设模型限流标记，避免设 29s 限流导致后续请求连续快速失败； + // 是否单账号分组在 antigravity 账号遇到 503 或模型限流预检查时才判定。 + c.Request = c.Request.WithContext(h.gatewayService.WithSingleAntigravityAccountGroupProbe(c.Request.Context(), apiKey.GroupID)) for { selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制 @@ -639,12 +640,9 @@ func (h *GatewayHandler) Messages(c *gin.Context) { } fallbackUsed := false - // 单账号分组提前设置 SingleAccountRetry 标记，让 Service 层首次 503 就不设模型限流标记。 - // 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流，导致后续请求连续快速失败。 - if h.gatewayService.IsSingleAntigravityAccountGroup(c.Request.Context(), currentAPIKey.GroupID) { - ctx := service.WithSingleAccountRetry(c.Request.Context(), true, h.metadataBridgeEnabled()) - c.Request = c.Request.WithContext(ctx) - } + // 单账号分组首次 503 就不设模型限流标记，避免设 29s 限流导致后续请求连续快速失败； + // 是否单账号分组在 antigravity 账号遇到 503 或模型限流预检查时才判定。 + c.Request = c.Request.WithContext(h.gatewayService.WithSingleAntigravityAccountGroupProbe(c.Request.Context(), currentAPIKey.GroupID)) // 会话槽失败释放：failover 链上每个选中的 Anthropic OAuth 账号都注册过同一会话 // （checkAndRegisterSession）。若请求最终失败（转发失败/客户端中断/换号耗尽）， diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index cc729c4c226e..73065e22ee35 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -113,6 +113,10 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { "This group is restricted to Claude Code clients (/v1/messages only)") return } + if geminiImageGenerationRejected(c, apiKey, nil, reqModel, channelMapping.MappedModel) { + h.chatCompletionsErrorResponse(c, http.StatusBadRequest, "invalid_request_error", service.ImageGenerationPermissionMessage()) + return + } if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, reqModel, body); decision != nil && !decision.AllowNextStage { h.openAISecurityAuditError(c, decision) @@ -317,6 +321,44 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { accountReleaseFunc() } + // 提交 usage 记录。成功路径与"流中断但 Forward 已观测到 usage 的部分结果" + // 错误路径共用；failover 错误恒定 result=nil，不会重复计费。 + submitForwardUsage := func(result *service.ForwardResult) { + userAgent := c.GetHeader("User-Agent") + clientIP := ip.GetClientIP(c) + requestPayloadHash := service.HashUsageRequestPayload(body) + inboundEndpoint := GetInboundEndpoint(c) + upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) + + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) + stampForwardRequestedReasoningEffort(result, service.RequestedReasoningEffortFromContext(c.Request.Context())) + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { + if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ + Result: result, + QuotaPlatform: quotaPlatform, + APIKey: apiKey, + User: apiKey.User, + Account: account, + Subscription: subscription, + PricingAt: pricingAt, + InboundEndpoint: inboundEndpoint, + UpstreamEndpoint: upstreamEndpoint, + UserAgent: userAgent, + IPAddress: clientIP, + RequestPayloadHash: requestPayloadHash, + APIKeyService: h.apiKeyService, + SessionID: sessionID, + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), + }); err != nil { + reqLog.Error("gateway.cc.record_usage_failed", + zap.Int64("account_id", account.ID), + zap.Error(err), + ) + } + }) + } + if err != nil { var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { @@ -347,43 +389,14 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { zap.Bool("upstream_error_response_already_written", upstreamErrorAlreadyCommunicated), zap.Error(err), ) + if result != nil { + submitForwardUsage(result) + } return } // 6. Record usage - userAgent := c.GetHeader("User-Agent") - clientIP := ip.GetClientIP(c) - requestPayloadHash := service.HashUsageRequestPayload(body) - inboundEndpoint := GetInboundEndpoint(c) - upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) - - quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - sessionID := service.ExtractClientSessionID(c) - stampForwardRequestedReasoningEffort(result, service.RequestedReasoningEffortFromContext(c.Request.Context())) - h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { - if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ - Result: result, - QuotaPlatform: quotaPlatform, - APIKey: apiKey, - User: apiKey.User, - Account: account, - Subscription: subscription, - PricingAt: pricingAt, - InboundEndpoint: inboundEndpoint, - UpstreamEndpoint: upstreamEndpoint, - UserAgent: userAgent, - IPAddress: clientIP, - RequestPayloadHash: requestPayloadHash, - APIKeyService: h.apiKeyService, - SessionID: sessionID, - ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), - }); err != nil { - reqLog.Error("gateway.cc.record_usage_failed", - zap.Int64("account_id", account.ID), - zap.Error(err), - ) - } - }) + submitForwardUsage(result) return } } diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index f54bb72edf35..90890a663f26 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -117,6 +117,10 @@ func (h *GatewayHandler) Responses(c *gin.Context) { "This group is restricted to Claude Code clients (/v1/messages only)") return } + if geminiImageGenerationRejected(c, apiKey, nil, reqModel, channelMapping.MappedModel) { + h.responsesErrorResponse(c, http.StatusBadRequest, "invalid_request_error", service.ImageGenerationPermissionMessage()) + return + } if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, body); decision != nil && !decision.AllowNextStage { h.responsesSecurityAuditError(c, decision) @@ -299,6 +303,44 @@ func (h *GatewayHandler) Responses(c *gin.Context) { accountReleaseFunc() } + // 提交 usage 记录。成功路径与"流中断但 Forward 已观测到 usage 的部分结果" + // 错误路径共用；failover 错误恒定 result=nil，不会重复计费。 + submitForwardUsage := func(result *service.ForwardResult) { + userAgent := c.GetHeader("User-Agent") + clientIP := ip.GetClientIP(c) + requestPayloadHash := service.HashUsageRequestPayload(body) + inboundEndpoint := GetInboundEndpoint(c) + upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) + + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + sessionID := service.ExtractClientSessionID(c) + stampForwardRequestedReasoningEffort(result, service.RequestedReasoningEffortFromContext(c.Request.Context())) + h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { + if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ + Result: result, + QuotaPlatform: quotaPlatform, + APIKey: apiKey, + User: apiKey.User, + Account: account, + Subscription: subscription, + PricingAt: pricingAt, + InboundEndpoint: inboundEndpoint, + UpstreamEndpoint: upstreamEndpoint, + UserAgent: userAgent, + IPAddress: clientIP, + RequestPayloadHash: requestPayloadHash, + APIKeyService: h.apiKeyService, + SessionID: sessionID, + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), + }); err != nil { + reqLog.Error("gateway.responses.record_usage_failed", + zap.Int64("account_id", account.ID), + zap.Error(err), + ) + } + }) + } + if err != nil { var failoverErr *service.UpstreamFailoverError if errors.As(err, &failoverErr) { @@ -330,43 +372,14 @@ func (h *GatewayHandler) Responses(c *gin.Context) { zap.Bool("upstream_error_response_already_written", upstreamErrorAlreadyCommunicated), zap.Error(err), ) + if result != nil { + submitForwardUsage(result) + } return } // 6. Record usage - userAgent := c.GetHeader("User-Agent") - clientIP := ip.GetClientIP(c) - requestPayloadHash := service.HashUsageRequestPayload(body) - inboundEndpoint := GetInboundEndpoint(c) - upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) - - quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) - sessionID := service.ExtractClientSessionID(c) - stampForwardRequestedReasoningEffort(result, service.RequestedReasoningEffortFromContext(c.Request.Context())) - h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) { - if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{ - Result: result, - QuotaPlatform: quotaPlatform, - APIKey: apiKey, - User: apiKey.User, - Account: account, - Subscription: subscription, - PricingAt: pricingAt, - InboundEndpoint: inboundEndpoint, - UpstreamEndpoint: upstreamEndpoint, - UserAgent: userAgent, - IPAddress: clientIP, - RequestPayloadHash: requestPayloadHash, - APIKeyService: h.apiKeyService, - SessionID: sessionID, - ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), - }); err != nil { - reqLog.Error("gateway.responses.record_usage_failed", - zap.Int64("account_id", account.ID), - zap.Error(err), - ) - } - }) + submitForwardUsage(result) return } } diff --git a/backend/internal/handler/gateway_handler_warmup_intercept_unit_test.go b/backend/internal/handler/gateway_handler_warmup_intercept_unit_test.go index aad6fcaad1f7..a4e37119f21a 100644 --- a/backend/internal/handler/gateway_handler_warmup_intercept_unit_test.go +++ b/backend/internal/handler/gateway_handler_warmup_intercept_unit_test.go @@ -158,8 +158,12 @@ func (f *fakeConcurrencyCache) CleanupStaleProcessSlots(context.Context, string) func newTestGatewayHandler(t *testing.T, group *service.Group, accounts []*service.Account) (*GatewayHandler, func()) { t.Helper() + return newTestGatewayHandlerWithSchedulerCache(t, group, &fakeSchedulerCache{accounts: accounts}) +} + +func newTestGatewayHandlerWithSchedulerCache(t *testing.T, group *service.Group, schedulerCache service.SchedulerCache) (*GatewayHandler, func()) { + t.Helper() - schedulerCache := &fakeSchedulerCache{accounts: accounts} schedulerSnapshot := service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil) gwSvc := service.NewGatewayService( diff --git a/backend/internal/handler/gemini_client_cancel_test.go b/backend/internal/handler/gemini_client_cancel_test.go index aad811f025b2..6fda17db5a9a 100644 --- a/backend/internal/handler/gemini_client_cancel_test.go +++ b/backend/internal/handler/gemini_client_cancel_test.go @@ -34,11 +34,12 @@ func (u *clientCancelUpstream) Do(req *http.Request, _ string, _ int64, _ int) ( } type geminiClientCancelFixture struct { - handler *GatewayHandler - group *service.Group - apiKey *service.APIKey - upstream *clientCancelUpstream - ctx context.Context + handler *GatewayHandler + group *service.Group + apiKey *service.APIKey + upstream *clientCancelUpstream + schedulerCache *bucketRecordingSchedulerCache + ctx context.Context } func newGeminiClientCancelFixture(t *testing.T) *geminiClientCancelFixture { @@ -57,7 +58,8 @@ func newGeminiClientCancelFixture(t *testing.T) *geminiClientCancelFixture { Priority: 1, AccountGroups: []service.AccountGroup{{AccountID: accountID, GroupID: groupID}}, } - h, cleanup := newTestGatewayHandler(t, group, []*service.Account{account}) + schedulerCache := &bucketRecordingSchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{account}}} + h, cleanup := newTestGatewayHandlerWithSchedulerCache(t, group, schedulerCache) t.Cleanup(cleanup) ctx, cancel := context.WithCancel(context.Background()) @@ -69,7 +71,7 @@ func newGeminiClientCancelFixture(t *testing.T) *geminiClientCancelFixture { ID: 9202, UserID: 9203, GroupID: &groupID, Group: group, Status: service.StatusActive, User: &service.User{ID: 9203, Concurrency: 10, Balance: 100}, } - return &geminiClientCancelFixture{handler: h, group: group, apiKey: apiKey, upstream: upstream, ctx: ctx} + return &geminiClientCancelFixture{handler: h, group: group, apiKey: apiKey, upstream: upstream, schedulerCache: schedulerCache, ctx: ctx} } // serve 经 ops 错误日志中间件执行请求，返回响应与 ops 队列中的条目数。 diff --git a/backend/internal/handler/gemini_image_generation_gate.go b/backend/internal/handler/gemini_image_generation_gate.go new file mode 100644 index 000000000000..f1667ee6d5cd --- /dev/null +++ b/backend/internal/handler/gemini_image_generation_gate.go @@ -0,0 +1,17 @@ +package handler + +import ( + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" +) + +// geminiImageGenerationRejected 在选号之前按客户端请求名、渠道映射名与请求体 +// 执行 Gemini 系分组的生图门禁。命中时把请求标记为本地功能门禁并返回 true， +// 调用方按入口协议写出 400；账号级模型映射之后的判定由转发层负责。 +func geminiImageGenerationRejected(c *gin.Context, apiKey *service.APIKey, body []byte, models ...string) bool { + if apiKey == nil || !service.GeminiImageGenerationDenied(apiKey.Group, body, models...) { + return false + } + service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) + return true +} diff --git a/backend/internal/handler/gemini_image_generation_gate_test.go b/backend/internal/handler/gemini_image_generation_gate_test.go new file mode 100644 index 000000000000..3bc893a7d155 --- /dev/null +++ b/backend/internal/handler/gemini_image_generation_gate_test.go @@ -0,0 +1,169 @@ +//go:build unit + +package handler + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type geminiImageGateEndpoint struct { + name string + path string + body string + params gin.Params + forcePlatform string + call func(*GatewayHandler, *gin.Context) +} + +func geminiImageGateEndpoints(model string) []geminiImageGateEndpoint { + return []geminiImageGateEndpoint{ + { + name: "messages", path: "/v1/messages", + body: `{"model":"` + model + `","max_tokens":16,"messages":[{"role":"user","content":"draw a cat"}]}`, + call: (*GatewayHandler).Messages, + }, + { + name: "chat completions", path: "/v1/chat/completions", + body: `{"model":"` + model + `","messages":[{"role":"user","content":"draw a cat"}]}`, + call: (*GatewayHandler).ChatCompletions, + }, + { + name: "responses", path: "/v1/responses", + body: `{"model":"` + model + `","input":"draw a cat"}`, + call: (*GatewayHandler).Responses, + }, + { + name: "gemini native", path: "/v1beta/models/" + model + ":generateContent", + body: `{"contents":[{"role":"user","parts":[{"text":"draw a cat"}]}]}`, + params: gin.Params{{Key: "modelAction", Value: "/" + model + ":generateContent"}}, + call: (*GatewayHandler).GeminiV1BetaModels, + }, + } +} + +func runGeminiImageGateEndpoint(t *testing.T, ep geminiImageGateEndpoint, group *service.Group, accountPlatform string) (*httptest.ResponseRecorder, *gin.Context, *groupScopedSchedulerCache) { + t.Helper() + const accountID = int64(9411) + schedulerCache := &groupScopedSchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{{ + ID: accountID, Platform: accountPlatform, Type: service.AccountTypeAPIKey, + Status: service.StatusActive, Schedulable: true, Concurrency: 1, + AccountGroups: []service.AccountGroup{{AccountID: accountID, GroupID: group.ID}}, + }}}} + gatewayService := service.NewGatewayService( + nil, &groupMapRepo{fakeGroupRepo: &fakeGroupRepo{}, groups: map[int64]*service.Group{group.ID: group}}, + nil, nil, nil, nil, nil, nil, nil, + service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil), + nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, + ) + cfg := &config.Config{RunMode: config.RunModeSimple} + billingCacheService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billingCacheService.Stop) + h := &GatewayHandler{ + gatewayService: gatewayService, + geminiCompatService: service.NewGeminiMessagesCompatService(nil, nil, nil, nil, nil, nil, nil, nil, cfg), + billingCacheService: billingCacheService, + concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatClaude, 0), + maxAccountSwitches: 1, + cfg: cfg, + } + + groupID := group.ID + apiKey := &service.APIKey{ + ID: 9420, UserID: 9430, GroupID: &groupID, Group: group, Status: service.StatusActive, + User: &service.User{ID: 9430, Concurrency: 10, Balance: 100}, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + ctx := context.WithValue(context.Background(), ctxkey.Group, group) + req := httptest.NewRequest(http.MethodPost, ep.path, bytes.NewBufferString(ep.body)).WithContext(ctx) + req.Header.Set("Content-Type", "application/json") + c.Request = req + c.Params = ep.params + if ep.forcePlatform != "" { + c.Set(string(middleware.ContextKeyForcePlatform), ep.forcePlatform) + } + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10}) + + ep.call(h, c) + return recorder, c, schedulerCache +} + +func TestGatewayHandlers_GeminiImageGenerationDisabledGroupRejectsBeforeSelection(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, platform := range []string{service.PlatformGemini, service.PlatformAntigravity} { + for _, ep := range geminiImageGateEndpoints("gemini-3-pro-image") { + if platform == service.PlatformAntigravity && ep.params != nil { + ep.path = "/antigravity" + ep.path + ep.forcePlatform = service.PlatformAntigravity + } + t.Run(platform+"/"+ep.name, func(t *testing.T) { + group := &service.Group{ + ID: 9400, Hydrated: true, Platform: platform, Status: service.StatusActive, + AllowImageGeneration: false, + } + recorder, c, schedulerCache := runGeminiImageGateEndpoint(t, ep, group, platform) + + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.Contains(t, recorder.Body.String(), service.ImageGenerationPermissionMessage()) + _, reachedSelection := c.Get(opsAccountIDKey) + require.False(t, reachedSelection) + require.Empty(t, schedulerCache.queriedGroupIDs()) + }) + } + } +} + +func TestGatewayHandlers_GeminiImageGenerationGateSkipsGroupsWithoutTheSwitch(t *testing.T) { + gin.SetMode(gin.TestMode) + // Anthropic 分组在后台不展示生图开关，混合调度 Antigravity 账号时不受门禁约束。 + for _, ep := range geminiImageGateEndpoints("gemini-3-pro-image")[:3] { + t.Run(ep.name, func(t *testing.T) { + group := &service.Group{ + ID: 9401, Hydrated: true, Platform: service.PlatformAnthropic, Status: service.StatusActive, + AllowImageGeneration: false, + } + recorder, _, schedulerCache := runGeminiImageGateEndpoint(t, ep, group, service.PlatformAnthropic) + + require.NotContains(t, recorder.Body.String(), service.ImageGenerationPermissionMessage()) + require.NotEmpty(t, schedulerCache.queriedGroupIDs()) + }) + } +} + +func TestGatewayHandlers_GeminiImageGenerationGateAllowsTextModelsAndEnabledGroups(t *testing.T) { + gin.SetMode(gin.TestMode) + cases := []struct { + name string + model string + allow bool + }{ + {name: "text model in disabled group", model: "gemini-3.7-flash", allow: false}, + {name: "image model in enabled group", model: "gemini-3-pro-image", allow: true}, + } + for _, tc := range cases { + for _, ep := range geminiImageGateEndpoints(tc.model)[:3] { + t.Run(tc.name+"/"+ep.name, func(t *testing.T) { + group := &service.Group{ + ID: 9402, Hydrated: true, Platform: service.PlatformGemini, Status: service.StatusActive, + AllowImageGeneration: tc.allow, + } + recorder, _, schedulerCache := runGeminiImageGateEndpoint(t, ep, group, service.PlatformGemini) + + require.NotContains(t, recorder.Body.String(), service.ImageGenerationPermissionMessage()) + require.NotEmpty(t, schedulerCache.queriedGroupIDs()) + }) + } + } +} diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 95c550302e6c..c15d212a3be2 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -328,6 +328,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { googleError(c, http.StatusBadRequest, "Invalid model in URL") return } + urlModel := modelName if resolvedModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()); ok && strings.TrimSpace(resolvedModel) != "" { modelName = strings.TrimSpace(resolvedModel) } @@ -358,14 +359,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context()) c.Request = c.Request.WithContext(pricingCtx) - if decision := h.checkSecurityAudit(c, reqLog, apiKey, authSubject, service.ContentModerationProtocolGemini, modelName, body); decision != nil && !decision.AllowNextStage { + // 解析渠道级模型映射 + channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, modelName) + reqModel := modelName // 保存映射前的原始模型名 + if action != "countTokens" && geminiImageGenerationRejected(c, apiKey, body, urlModel, reqModel, channelMapping.MappedModel) { + googleError(c, http.StatusBadRequest, service.ImageGenerationPermissionMessage()) + return + } + + if decision := h.checkSecurityAudit(c, reqLog, apiKey, authSubject, service.ContentModerationProtocolGemini, reqModel, body); decision != nil && !decision.AllowNextStage { googleSecurityAuditError(c, decision) return } - // 解析渠道级模型映射 - channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, modelName) - reqModel := modelName // 保存映射前的原始模型名 if channelMapping.Mapped { modelName = channelMapping.MappedModel } @@ -546,12 +552,9 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { } fs := NewFailoverState(maxAccountSwitches, hasBoundSession) - // 单账号分组提前设置 SingleAccountRetry 标记，让 Service 层首次 503 就不设模型限流标记。 - // 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流，导致后续请求连续快速失败。 - if h.gatewayService.IsSingleAntigravityAccountGroup(c.Request.Context(), apiKey.GroupID) { - ctx := service.WithSingleAccountRetry(c.Request.Context(), true, h.metadataBridgeEnabled()) - c.Request = c.Request.WithContext(ctx) - } + // 单账号分组首次 503 就不设模型限流标记，避免设 29s 限流导致后续请求连续快速失败； + // 是否单账号分组在 antigravity 账号遇到 503 或模型限流预检查时才判定。 + c.Request = c.Request.WithContext(h.gatewayService.WithSingleAntigravityAccountGroupProbe(c.Request.Context(), apiKey.GroupID)) for { excludedAccountIDs := fs.FailedAccountIDs diff --git a/backend/internal/handler/single_account_retry_probe_unit_test.go b/backend/internal/handler/single_account_retry_probe_unit_test.go new file mode 100644 index 000000000000..8b73067e254d --- /dev/null +++ b/backend/internal/handler/single_account_retry_probe_unit_test.go @@ -0,0 +1,120 @@ +//go:build unit + +package handler + +import ( + "bytes" + "context" + "net/http/httptest" + "sync" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// bucketRecordingSchedulerCache 记录读取过的快照桶平台。 +type bucketRecordingSchedulerCache struct { + *fakeSchedulerCache + mu sync.Mutex + platforms []string +} + +func (c *bucketRecordingSchedulerCache) GetSnapshot(ctx context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) { + c.mu.Lock() + c.platforms = append(c.platforms, bucket.Platform) + c.mu.Unlock() + return c.fakeSchedulerCache.GetSnapshot(ctx, bucket) +} + +func (c *bucketRecordingSchedulerCache) readPlatforms() []string { + c.mu.Lock() + defer c.mu.Unlock() + return append([]string(nil), c.platforms...) +} + +// 非 antigravity 账号的请求不读取单 antigravity 账号分组判定，选号之外不查询 antigravity 快照桶。 +func TestGeminiGroupRequestsDoNotReadAntigravityBucket(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + route string + path string + body string + call func(h *GatewayHandler) func(*gin.Context) + }{ + { + name: "gemini native", + route: "/v1beta/models/*modelAction", + path: "/v1beta/models/gemini-2.5-flash:generateContent", + body: `{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`, + call: func(h *GatewayHandler) func(*gin.Context) { return h.GeminiV1BetaModels }, + }, + { + name: "messages on gemini group", + route: "/v1/messages", + path: "/v1/messages", + body: `{"model":"gemini-2.5-flash","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}`, + call: func(h *GatewayHandler) func(*gin.Context) { return h.Messages }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + f := newGeminiClientCancelFixture(t) + + f.serve(t, tt.route, tt.path, tt.body, tt.call(f.handler)) + + require.Equal(t, int32(1), f.upstream.calls.Load(), "请求应完成选号并到达上游") + platforms := f.schedulerCache.readPlatforms() + require.NotEmpty(t, platforms) + require.NotContains(t, platforms, service.PlatformAntigravity) + }) + } +} + +func TestMessagesAnthropicGroupDoesNotReadAntigravityBucket(t *testing.T) { + gin.SetMode(gin.TestMode) + groupID := int64(9400) + accountID := int64(9401) + group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformAnthropic, Status: service.StatusActive} + account := &service.Account{ + ID: accountID, + Platform: service.PlatformAnthropic, + Type: service.AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "test-key", "intercept_warmup_requests": true}, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 1, + AccountGroups: []service.AccountGroup{{AccountID: accountID, GroupID: groupID}}, + } + schedulerCache := &bucketRecordingSchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{account}}} + h, cleanup := newTestGatewayHandlerWithSchedulerCache(t, group, schedulerCache) + t.Cleanup(cleanup) + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := `{"model":"claude-sonnet-4-5","max_tokens":256,"messages":[{"role":"user","content":[{"type":"text","text":"Warmup"}]}]}` + req := httptest.NewRequest("POST", "/v1/messages", bytes.NewBufferString(body)) + req.Header.Set("Content-Type", "application/json") + c.Request = req.WithContext(context.WithValue(req.Context(), ctxkey.Group, group)) + apiKey := &service.APIKey{ + ID: 9402, UserID: 9403, GroupID: &groupID, Group: group, Status: service.StatusActive, + User: &service.User{ID: 9403, Concurrency: 10, Balance: 100}, + } + c.Set(string(middleware.ContextKeyAPIKey), apiKey) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10}) + + h.Messages(c) + + require.Equal(t, 200, rec.Code) + selected, ok := c.Get(opsAccountIDKey) + require.True(t, ok) + require.Equal(t, accountID, selected) + platforms := schedulerCache.readPlatforms() + require.NotEmpty(t, platforms) + require.NotContains(t, platforms, service.PlatformAntigravity) +} diff --git a/backend/internal/pkg/antigravity/response_transformer.go b/backend/internal/pkg/antigravity/response_transformer.go index 0db6822f315b..03e26ac50cd9 100644 --- a/backend/internal/pkg/antigravity/response_transformer.go +++ b/backend/internal/pkg/antigravity/response_transformer.go @@ -277,11 +277,12 @@ func (p *NonStreamingProcessor) buildResponse(geminiResp *GeminiResponse, respon } // 注意：Gemini 的 promptTokenCount 包含 cachedContentTokenCount， - // 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去 + // 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去。 + // prompt 被拦截时上游可能只报 cachedContentTokenCount，差值下限为 0。 usage := ClaudeUsage{} if geminiResp.UsageMetadata != nil { cached := geminiResp.UsageMetadata.CachedContentTokenCount - usage.InputTokens = geminiResp.UsageMetadata.PromptTokenCount - cached + usage.InputTokens = max(geminiResp.UsageMetadata.PromptTokenCount-cached, 0) usage.OutputTokens = geminiResp.UsageMetadata.CandidatesTokenCount + geminiResp.UsageMetadata.ThoughtsTokenCount usage.CacheReadInputTokens = cached usage.ImageOutputTokens = geminiResp.UsageMetadata.ImageOutputTokens() diff --git a/backend/internal/pkg/antigravity/stream_transformer.go b/backend/internal/pkg/antigravity/stream_transformer.go index faecf233ecac..6b52ba882f93 100644 --- a/backend/internal/pkg/antigravity/stream_transformer.go +++ b/backend/internal/pkg/antigravity/stream_transformer.go @@ -109,10 +109,11 @@ func (p *StreamingProcessor) ProcessLine(line string) []byte { // 更新 usage // 注意：Gemini 的 promptTokenCount 包含 cachedContentTokenCount， - // 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去 + // 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去。 + // prompt 被拦截时上游可能只报 cachedContentTokenCount，差值下限为 0。 if geminiResp.UsageMetadata != nil { cached := geminiResp.UsageMetadata.CachedContentTokenCount - p.inputTokens = geminiResp.UsageMetadata.PromptTokenCount - cached + p.inputTokens = max(geminiResp.UsageMetadata.PromptTokenCount-cached, 0) p.outputTokens = geminiResp.UsageMetadata.CandidatesTokenCount + geminiResp.UsageMetadata.ThoughtsTokenCount p.cacheReadTokens = cached p.imageOutputTokens = geminiResp.UsageMetadata.ImageOutputTokens() @@ -190,7 +191,7 @@ func (p *StreamingProcessor) emitMessageStart(v1Resp *V1InternalResponse) []byte usage := ClaudeUsage{} if v1Resp.Response.UsageMetadata != nil { cached := v1Resp.Response.UsageMetadata.CachedContentTokenCount - usage.InputTokens = v1Resp.Response.UsageMetadata.PromptTokenCount - cached + usage.InputTokens = max(v1Resp.Response.UsageMetadata.PromptTokenCount-cached, 0) usage.OutputTokens = v1Resp.Response.UsageMetadata.CandidatesTokenCount + v1Resp.Response.UsageMetadata.ThoughtsTokenCount usage.CacheReadInputTokens = cached usage.ImageOutputTokens = v1Resp.Response.UsageMetadata.ImageOutputTokens() diff --git a/backend/internal/pkg/antigravity/usage_contract_test.go b/backend/internal/pkg/antigravity/usage_contract_test.go index 832bffbfb2f7..e52431742f2f 100644 --- a/backend/internal/pkg/antigravity/usage_contract_test.go +++ b/backend/internal/pkg/antigravity/usage_contract_test.go @@ -38,3 +38,29 @@ func TestGeminiUsageMapping_NoCacheCreationTokens(t *testing.T) { require.Zero(t, usage.CacheCreationInputTokens) }) } + +// usageMetadata 缺 promptTokenCount、只有 cachedContentTokenCount 时，input_tokens 下限为 0， +// 缓存命中仍计入 cache_read_input_tokens。 +func TestGeminiUsageMapping_MissingPromptTokenCountKeepsInputNonNegative(t *testing.T) { + const geminiBody = `{"candidates":[{"content":{"parts":[{"text":"hi"}],"role":"model"},"finishReason":"STOP"}],` + + `"usageMetadata":{"cachedContentTokenCount":30}}` + + t.Run("non-stream", func(t *testing.T) { + _, usage, err := TransformGeminiToClaude([]byte(geminiBody), "gemini-3.1-pro-preview") + require.NoError(t, err) + require.NotNil(t, usage) + require.Zero(t, usage.InputTokens) + require.Equal(t, 30, usage.CacheReadInputTokens) + }) + + t.Run("stream", func(t *testing.T) { + p := NewStreamingProcessor("gemini-3.1-pro-preview") + out := p.ProcessLine(`data: {"response":` + geminiBody + `}`) + require.True(t, strings.Contains(string(out), `"message_start"`)) + require.NotContains(t, string(out), `"input_tokens":-`, "message_start 的 input_tokens 不应为负") + _, usage := p.Finish() + require.NotNil(t, usage) + require.Zero(t, usage.InputTokens) + require.Equal(t, 30, usage.CacheReadInputTokens) + }) +} diff --git a/backend/internal/pkg/apicompat/responses_chat_protocol_completion_test.go b/backend/internal/pkg/apicompat/responses_chat_protocol_completion_test.go deleted file mode 100644 index 646963874053..000000000000 --- a/backend/internal/pkg/apicompat/responses_chat_protocol_completion_test.go +++ /dev/null @@ -1,83 +0,0 @@ -package apicompat - -import ( - "testing" - - "github.com/stretchr/testify/require" -) - -func TestResponsesChatFinalizeDoesNotCompleteProtocolFailure(t *testing.T) { - for _, failure := range []string{"arguments", "call_id", "missing_terminal_call"} { - t.Run(failure, func(t *testing.T) { - state := NewResponsesEventToChatState() - state.IncludeUsage = true - item := ResponsesOutput{Type: "function_call", ID: "fc_fixture", CallID: "call_fixture", Name: "lookup", Arguments: `{"x":1}`} - require.NotEmpty(t, ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.output_item.added", OutputIndex: 3, Item: &item}, state)) - end := item - var failureEvent ResponsesStreamEvent - switch failure { - case "arguments": - end.Arguments = `{"x":2}` - failureEvent = ResponsesStreamEvent{Type: "response.output_item.done", OutputIndex: 3, Item: &end} - case "call_id": - end.CallID = "call_changed" - failureEvent = ResponsesStreamEvent{Type: "response.output_item.done", OutputIndex: 3, Item: &end} - case "missing_terminal_call": - failureEvent = ResponsesStreamEvent{Type: "response.completed", Response: &ResponsesResponse{ - Status: "completed", Output: []ResponsesOutput{{Type: "message", Role: "assistant"}}, - Usage: &ResponsesUsage{InputTokens: 13, OutputTokens: 5}, - }} - } - require.Empty(t, ResponsesEventToChatChunks(&failureEvent, state)) - require.NotEmpty(t, state.ProtocolError, "the real converter must detect the contradiction") - require.Empty(t, FinalizeResponsesChatStream(state), "a failed stream must never acquire a normal finish or usage chunk") - require.Empty(t, ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.completed", Response: &ResponsesResponse{ - Status: "completed", Output: []ResponsesOutput{item}, Usage: &ResponsesUsage{InputTokens: 13, OutputTokens: 5}, - }}, state), "a late terminal event cannot recover a corrupted stream") - require.Empty(t, FinalizeResponsesChatStream(state), "repeated finalization must remain silent") - }) - } -} - -func TestResponsesChatProtocolCompletionPreservesValidToolAndUsage(t *testing.T) { - state := NewResponsesEventToChatState() - state.IncludeUsage = true - initial := ResponsesOutput{Type: "function_call", ID: "fc_fixture", CallID: "call_fixture", Name: "lookup"} - complete := initial - complete.Arguments, complete.Status = `{"x":1}`, "completed" - events := []ResponsesStreamEvent{ - {Type: "response.output_item.added", OutputIndex: 3, Item: &initial}, - {Type: "response.function_call_arguments.delta", OutputIndex: 3, Delta: `{"x":`}, - {Type: "response.output_item.done", OutputIndex: 3, Item: &complete}, - {Type: "response.completed", Response: &ResponsesResponse{Status: "completed", Output: []ResponsesOutput{complete}, Usage: &ResponsesUsage{InputTokens: 13, OutputTokens: 5}}}, - } - var args, callID, name string - var finishes, usages int - for i := range events { - for _, chunk := range ResponsesEventToChatChunks(&events[i], state) { - if chunk.Usage != nil { - usages++ - require.Equal(t, 13, chunk.Usage.PromptTokens) - require.Equal(t, 5, chunk.Usage.CompletionTokens) - } - for _, choice := range chunk.Choices { - if choice.FinishReason != nil { - finishes++ - require.Equal(t, "tool_calls", *choice.FinishReason) - } - for _, tool := range choice.Delta.ToolCalls { - args += tool.Function.Arguments - callID += tool.ID - name += tool.Function.Name - } - } - } - } - require.Empty(t, state.ProtocolError) - require.Equal(t, complete.Arguments, args) - require.Equal(t, complete.CallID, callID) - require.Equal(t, complete.Name, name) - require.Equal(t, 1, finishes) - require.Equal(t, 1, usages) - require.Empty(t, FinalizeResponsesChatStream(state)) -} diff --git a/backend/internal/pkg/apicompat/responses_to_chat_function_state.go b/backend/internal/pkg/apicompat/responses_to_chat_function_state.go deleted file mode 100644 index 08969690e918..000000000000 --- a/backend/internal/pkg/apicompat/responses_to_chat_function_state.go +++ /dev/null @@ -1,160 +0,0 @@ -package apicompat - -import ( - "bytes" - "crypto/sha256" - "encoding/json" - "hash" -) - -// The client may have already received a prefix of the function arguments. -// Keep only its length and digest: an authoritative snapshot may append a -// missing suffix, but cannot rewrite emitted bytes or require an unbounded copy. -type responsesChatFunctionState struct { - chatIndex int - itemID, callID, name string - argumentBytes int - argumentHash hash.Hash - typeSent, completed bool - argumentsDone bool // a complete argument-done snapshot was observed -} - -func responseChatFunctionError(state *ResponsesEventToChatState) { - state.ProtocolError = "Upstream Responses function-call stream is incomplete or inconsistent" -} - -func resToChatFunctionItem(item *ResponsesOutput, outputIndex int, complete bool, state *ResponsesEventToChatState) []ChatCompletionsChunk { - if item == nil || item.Type != "function_call" || outputIndex < 0 || state.ProtocolError != "" { - return nil - } - if state.functionCalls == nil { - state.functionCalls = make(map[int]*responsesChatFunctionState) - } - if state.OutputIndexToToolIndex == nil { - state.OutputIndexToToolIndex = make(map[int]int) - } - tool := state.functionCalls[outputIndex] - if tool == nil { - if _, occupied := state.OutputIndexToToolIndex[outputIndex]; occupied { - responseChatFunctionError(state) - return nil - } - tool = &responsesChatFunctionState{chatIndex: state.NextToolCallIndex, argumentHash: sha256.New()} - state.NextToolCallIndex++ - state.functionCalls[outputIndex] = tool - state.OutputIndexToToolIndex[outputIndex] = tool.chatIndex - } - if (tool.itemID != "" && item.ID != "" && tool.itemID != item.ID) || - (tool.callID != "" && item.CallID != "" && tool.callID != item.CallID) || - (tool.name != "" && item.Name != "" && tool.name != item.Name) { - responseChatFunctionError(state) - return nil - } - if item.CallID != "" { - for index, other := range state.functionCalls { - if index != outputIndex && other.callID == item.CallID { - responseChatFunctionError(state) - return nil - } - } - } - tool.itemID = firstNonemptyToolIdentity(tool.itemID, item.ID) - delta := ChatToolCall{Index: &tool.chatIndex} - if !tool.typeSent { - delta.Type = "function" - tool.typeSent = true - } - if tool.callID == "" && item.CallID != "" { - tool.callID, delta.ID = item.CallID, item.CallID - } - if tool.name == "" && item.Name != "" { - tool.name, delta.Function.Name = item.Name, item.Name - } - hasCompleteArguments := json.Valid([]byte(item.Arguments)) || (item.Arguments == "" && tool.argumentsDone) - if complete && (tool.callID == "" || tool.name == "" || !hasCompleteArguments || - (item.Status != "" && item.Status != "completed")) { - responseChatFunctionError(state) - return nil - } - if item.Arguments != "" || (complete && !tool.argumentsDone) { - var ok bool - delta.Function.Arguments, ok = resToChatArgumentsSnapshotSuffix(item.Arguments, tool, state) - if !ok { - return nil - } - } - tool.completed = tool.completed || complete - state.SawToolCall = true - if delta.Type == "" && delta.ID == "" && delta.Function.Name == "" && delta.Function.Arguments == "" { - return nil - } - chatDelta := ChatDelta{ToolCalls: []ChatToolCall{delta}} - if !state.SentRole { - chatDelta.Role = "assistant" - state.SentRole = true - } - return []ChatCompletionsChunk{makeChatDeltaChunk(state, chatDelta)} -} - -func resToChatArgumentsSnapshotSuffix(arguments string, tool *responsesChatFunctionState, state *ResponsesEventToChatState) (string, bool) { - if tool.argumentBytes > len(arguments) || ((tool.completed || tool.argumentsDone) && tool.argumentBytes != len(arguments)) { - responseChatFunctionError(state) - return "", false - } - prefix := sha256.Sum256([]byte(arguments[:tool.argumentBytes])) - if !bytes.Equal(prefix[:], tool.argumentHash.Sum(nil)) { - responseChatFunctionError(state) - return "", false - } - suffix := arguments[tool.argumentBytes:] - _, _ = tool.argumentHash.Write([]byte(suffix)) - tool.argumentBytes = len(arguments) - return suffix, true -} - -func firstNonemptyToolIdentity(existing, incoming string) string { - if existing != "" { - return existing - } - return incoming -} - -func resToChatCompleteFunctions(response *ResponsesResponse, state *ResponsesEventToChatState) []ChatCompletionsChunk { - if response == nil || response.Status != "completed" { - return nil - } - if !state.SentRole && !state.SawToolCall && !state.SawText && response.ID != "" { - state.ID = response.ID - } - var chunks []ChatCompletionsChunk - seen := make(map[int]bool) - for position := range response.Output { - item := &response.Output[position] - if item.Type != "function_call" { - continue - } - index := position - // The final output array need not repeat a sparse source output_index. - // Use actual call/item identity, never a guessed tool name or call ID. - for candidate, tool := range state.functionCalls { - if (item.CallID != "" && tool.callID == item.CallID) || (item.ID != "" && tool.itemID == item.ID) { - index = candidate - break - } - } - chunks = append(chunks, resToChatFunctionItem(item, index, true, state)...) - seen[index] = true - if state.ProtocolError != "" { - return nil - } - } - if len(response.Output) > 0 { - for index := range state.functionCalls { - if !seen[index] { - responseChatFunctionError(state) - return nil - } - } - } - return chunks -} diff --git a/backend/internal/pkg/apicompat/responses_to_chat_function_state_test.go b/backend/internal/pkg/apicompat/responses_to_chat_function_state_test.go index 42e7ad09a6e4..f047181deda1 100644 --- a/backend/internal/pkg/apicompat/responses_to_chat_function_state_test.go +++ b/backend/internal/pkg/apicompat/responses_to_chat_function_state_test.go @@ -6,105 +6,6 @@ import ( "github.com/stretchr/testify/require" ) -func TestResponsesEventToChatFunctionSnapshots(t *testing.T) { - final := ResponsesOutput{Type: "function_call", ID: "fc_one", CallID: "call_one", Name: "load_orders", Arguments: `{"city":"杭州"}`, Status: "completed"} - initial := final - initial.Arguments, initial.Status = "", "in_progress" - created := ResponsesStreamEvent{Type: "response.created", Response: &ResponsesResponse{ID: "resp_one"}} - added := ResponsesStreamEvent{Type: "response.output_item.added", OutputIndex: 7, Item: &initial} - done := ResponsesStreamEvent{Type: "response.output_item.done", OutputIndex: 7, Item: &final} - completed := ResponsesStreamEvent{Type: "response.completed", Response: &ResponsesResponse{ID: "resp_one", Status: "completed", Output: []ResponsesOutput{final}}} - argumentsDone := ResponsesStreamEvent{Type: "response.function_call_arguments.done", OutputIndex: 7, ItemID: final.ID, CallID: final.CallID, Name: final.Name, Arguments: final.Arguments} - emptyFinal := final - emptyFinal.Arguments = "" - emptyDone := ResponsesStreamEvent{Type: "response.output_item.done", OutputIndex: 7, Item: &emptyFinal} - emptyCompleted := ResponsesStreamEvent{Type: "response.completed", Response: &ResponsesResponse{ID: "resp_one", Status: "completed", Output: []ResponsesOutput{emptyFinal}}} - for name, events := range map[string][]ResponsesStreamEvent{ - "added_full": {created, {Type: "response.output_item.added", OutputIndex: 7, Item: &final}, done, completed}, - "done_full": {created, added, done, completed}, - "terminal_full": {created, added, completed}, - "terminal_only": {completed}, - "done_only": {created, done, completed}, - "partial_suffix": {created, added, {Type: "response.function_call_arguments.delta", OutputIndex: 7, Delta: `{"city":`}, done, completed}, - "arguments_done": {created, added, argumentsDone, done, completed}, - "arguments_done_partial_suffix": {created, added, {Type: "response.function_call_arguments.delta", OutputIndex: 7, Delta: `{"city":`}, argumentsDone, done, completed}, - "arguments_done_repeated": {created, added, argumentsDone, argumentsDone, done, argumentsDone, completed}, - "arguments_done_empty_terminal_snapshots": {created, added, argumentsDone, emptyDone, emptyCompleted}, - } { - t.Run(name, func(t *testing.T) { - state := NewResponsesEventToChatState() - state.Model = "fixture" - var arguments, callID, toolName, finish, role string - for i := range events { - chunks := ResponsesEventToChatChunks(&events[i], state) - for _, chunk := range chunks { - require.Equal(t, "resp_one", chunk.ID) - for _, choice := range chunk.Choices { - role += choice.Delta.Role - for _, call := range choice.Delta.ToolCalls { - require.NotNil(t, call.Index) - require.Equal(t, 0, *call.Index) - arguments += call.Function.Arguments - callID += call.ID - toolName += call.Function.Name - } - if choice.FinishReason != nil { - finish = *choice.FinishReason - } - } - } - } - require.Empty(t, state.ProtocolError) - require.Equal(t, final.Arguments, arguments, "the complete snapshot fills only the missing suffix") - require.Equal(t, final.CallID, callID, "never duplicate or invent call identity") - require.Equal(t, final.Name, toolName) - require.Equal(t, "assistant", role) - require.Equal(t, "tool_calls", finish) - }) - } -} - -func TestResponsesEventToChatArgumentsDoneRejectsContradictingPrefix(t *testing.T) { - for _, custom := range []bool{false, true} { - t.Run(map[bool]string{false: "function", true: "custom"}[custom], func(t *testing.T) { - state := NewResponsesEventToChatState() - itemType, deltaType, doneType := "function_call", "response.function_call_arguments.delta", "response.function_call_arguments.done" - if custom { - itemType, deltaType, doneType = "custom_tool_call", "response.custom_tool_call_input.delta", "response.custom_tool_call_input.done" - } - ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.output_item.added", OutputIndex: 3, Item: &ResponsesOutput{Type: itemType, ID: "item", CallID: "call", Name: "load"}}, state) - ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: deltaType, OutputIndex: 3, Delta: `{"x":1`}, state) - done := ResponsesStreamEvent{Type: doneType, OutputIndex: 3, Arguments: `{"x":2}`, Input: `{"x":2}`} - require.Nil(t, ResponsesEventToChatChunks(&done, state)) - require.NotEmpty(t, state.ProtocolError) - }) - } -} - -func TestResponsesEventToChatCustomInputDoneSharesPrefixState(t *testing.T) { - state := NewResponsesEventToChatState() - events := []ResponsesStreamEvent{ - {Type: "response.output_item.added", OutputIndex: 4, Item: &ResponsesOutput{Type: "custom_tool_call", CallID: "call_patch", Name: "apply_patch"}}, - {Type: "response.custom_tool_call_input.delta", OutputIndex: 4, Delta: "*** Begin"}, - {Type: "response.custom_tool_call_input.done", OutputIndex: 4, Input: "*** Begin Patch\n*** End Patch"}, - {Type: "response.custom_tool_call_input.done", OutputIndex: 4, Input: "*** Begin Patch\n*** End Patch"}, - } - var input string - for i := range events { - for _, chunk := range ResponsesEventToChatChunks(&events[i], state) { - for _, choice := range chunk.Choices { - for _, call := range choice.Delta.ToolCalls { - input += call.Function.Arguments - } - } - } - } - require.Empty(t, state.ProtocolError) - require.Equal(t, "*** Begin Patch\n*** End Patch", input) - require.Nil(t, ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.custom_tool_call_input.delta", OutputIndex: 4, Delta: "late"}, state)) - require.NotEmpty(t, state.ProtocolError, "an authoritative done must not accept a second tail") -} - func TestBufferedResponseAccumulator_ArgumentDonePreservesCallIdentity(t *testing.T) { acc := NewBufferedResponseAccumulator() acc.ProcessEvent(&ResponsesStreamEvent{Type: "response.output_item.added", OutputIndex: 0, Item: &ResponsesOutput{Type: "function_call", CallID: "call_original", Name: "lookup"}}) @@ -133,51 +34,3 @@ func TestBufferedResponseAccumulator_PartialDeltaDoesNotFillTerminalArguments(t acc.SupplementResponseOutput(resp) require.Empty(t, resp.Output[0].Arguments) } - -func TestResponsesEventToChatFunctionSnapshotRejectsContradiction(t *testing.T) { - for _, change := range []string{"arguments", "call_id", "name", "malformed", "after_done"} { - t.Run(change, func(t *testing.T) { - state := NewResponsesEventToChatState() - item := ResponsesOutput{Type: "function_call", ID: "fc", CallID: "call", Name: "load", Arguments: `{"x":1}`} - ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.output_item.added", Item: &item}, state) - end := item - switch change { - case "arguments": - end.Arguments = `{"x":2}` - case "call_id": - end.CallID = "different" - case "name": - end.Name = "different" - case "malformed": - end.Arguments = `{` - } - ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.output_item.done", Item: &end}, state) - if change == "after_done" { - ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.function_call_arguments.delta", Delta: "extra"}, state) - } - require.NotEmpty(t, state.ProtocolError) - require.Nil(t, ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.completed", Response: &ResponsesResponse{Status: "completed"}}, state)) - }) - } -} - -func TestResponsesEventToChatFunctionSnapshotLateIdentity(t *testing.T) { - state := NewResponsesEventToChatState() - ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.output_item.added", OutputIndex: 2, Item: &ResponsesOutput{Type: "function_call", ID: "fc"}}, state) - ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.function_call_arguments.delta", OutputIndex: 2, Delta: "{"}, state) - chunks := ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.output_item.done", OutputIndex: 2, Item: &ResponsesOutput{Type: "function_call", ID: "fc", CallID: "call_real", Name: "load_orders", Arguments: "{}"}}, state) - require.Empty(t, state.ProtocolError) - require.Len(t, chunks, 1) - call := chunks[0].Choices[0].Delta.ToolCalls[0] - require.Equal(t, "call_real", call.ID) - require.Equal(t, "load_orders", call.Function.Name) - require.Equal(t, "}", call.Function.Arguments) -} - -func TestResponsesEventToChatFunctionSnapshotCannotDropAnnouncedCall(t *testing.T) { - state := NewResponsesEventToChatState() - ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.output_item.added", OutputIndex: 7, Item: &ResponsesOutput{Type: "function_call", CallID: "call_original", Name: "load", Arguments: "{}"}}, state) - chunks := ResponsesEventToChatChunks(&ResponsesStreamEvent{Type: "response.completed", Response: &ResponsesResponse{Status: "completed", Output: []ResponsesOutput{{Type: "function_call", CallID: "call_other", Name: "load", Arguments: "{}"}}}}, state) - require.NotEmpty(t, state.ProtocolError) - require.Nil(t, chunks) -} diff --git a/backend/internal/pkg/apicompat/responses_to_chatcompletions.go b/backend/internal/pkg/apicompat/responses_to_chatcompletions.go index d1a7c5d5cee7..0b4daaa8d5f8 100644 --- a/backend/internal/pkg/apicompat/responses_to_chatcompletions.go +++ b/backend/internal/pkg/apicompat/responses_to_chatcompletions.go @@ -2,7 +2,6 @@ package apicompat import ( "crypto/rand" - "crypto/sha256" "encoding/hex" "encoding/json" "fmt" @@ -132,11 +131,9 @@ type ResponsesEventToChatState struct { Finalized bool // true after finish chunk has been emitted NextToolCallIndex int // next sequential tool_call index to assign OutputIndexToToolIndex map[int]int // Responses output_index → Chat tool_calls index + OutputIndexToArguments map[int]string IncludeUsage bool Usage *ChatUsage - ProtocolError string // safe static error; never embeds upstream arguments - functionCalls map[int]*responsesChatFunctionState - customToolArguments map[int]*responsesChatFunctionState } // NewResponsesEventToChatState returns an initialised stream state. @@ -145,15 +142,13 @@ func NewResponsesEventToChatState() *ResponsesEventToChatState { ID: generateChatCmplID(), Created: time.Now().Unix(), OutputIndexToToolIndex: make(map[int]int), + OutputIndexToArguments: make(map[int]string), } } // ResponsesEventToChatChunks converts a single Responses SSE event into zero // or more Chat Completions chunks, updating state as it goes. func ResponsesEventToChatChunks(evt *ResponsesStreamEvent, state *ResponsesEventToChatState) []ChatCompletionsChunk { - if state.ProtocolError != "" { - return nil - } switch evt.Type { case "response.created": return resToChatHandleCreated(evt, state) @@ -163,8 +158,6 @@ func ResponsesEventToChatChunks(evt *ResponsesStreamEvent, state *ResponsesEvent return resToChatHandleRefusalDelta(evt, state) case "response.output_item.added": return resToChatHandleOutputItemAdded(evt, state) - case "response.output_item.done": - return resToChatFunctionItem(evt.Item, evt.OutputIndex, true, state) case "response.function_call_arguments.delta", // custom/freeform 工具（如新版 apply_patch）的输入增量与 function_call 参数增量同形， // 均按 OutputIndex 累加到对应工具调用。 @@ -191,10 +184,9 @@ func ResponsesEventToChatChunks(evt *ResponsesStreamEvent, state *ResponsesEvent // FinalizeResponsesChatStream emits a final chunk with finish_reason if the // stream ended without a proper completion event (e.g. upstream disconnect). // It is idempotent: if a completion event already emitted the finish chunk, -// this returns nil. A protocol failure must be surfaced by the caller and -// cannot be converted into a normal finish or usage chunk at EOF. +// this returns nil. func FinalizeResponsesChatStream(state *ResponsesEventToChatState) []ChatCompletionsChunk { - if state.Finalized || state.ProtocolError != "" { + if state.Finalized { return nil } state.Finalized = true @@ -272,9 +264,6 @@ func resToChatHandleRefusalDelta(evt *ResponsesStreamEvent, state *ResponsesEven } func resToChatHandleOutputItemAdded(evt *ResponsesStreamEvent, state *ResponsesEventToChatState) []ChatCompletionsChunk { - if evt.Item != nil && evt.Item.Type == "function_call" { - return resToChatFunctionItem(evt.Item, evt.OutputIndex, false, state) - } // function_call 与 custom_tool_call（custom/freeform 工具）均按工具调用注册， // 以便后续 *_input.delta / *_arguments.delta 能映射到正确的工具索引。 if evt.Item == nil || (evt.Item.Type != "function_call" && evt.Item.Type != "custom_tool_call") { @@ -285,10 +274,6 @@ func resToChatHandleOutputItemAdded(evt *ResponsesStreamEvent, state *ResponsesE idx := state.NextToolCallIndex state.OutputIndexToToolIndex[evt.OutputIndex] = idx state.NextToolCallIndex++ - if state.customToolArguments == nil { - state.customToolArguments = make(map[int]*responsesChatFunctionState) - } - state.customToolArguments[evt.OutputIndex] = &responsesChatFunctionState{argumentHash: sha256.New()} return []ChatCompletionsChunk{makeChatDeltaChunk(state, ChatDelta{ ToolCalls: []ChatToolCall{{ @@ -311,18 +296,7 @@ func resToChatHandleFuncArgsDelta(evt *ResponsesStreamEvent, state *ResponsesEve if !ok { return nil } - tool := state.functionCalls[evt.OutputIndex] - if tool == nil { - tool = state.customToolArguments[evt.OutputIndex] - } - if tool != nil { - if tool.completed || tool.argumentsDone { - responseChatFunctionError(state) - return nil - } - _, _ = tool.argumentHash.Write([]byte(evt.Delta)) - tool.argumentBytes += len(evt.Delta) - } + state.OutputIndexToArguments[evt.OutputIndex] += evt.Delta return []ChatCompletionsChunk{makeChatDeltaChunk(state, ChatDelta{ ToolCalls: []ChatToolCall{{ @@ -340,34 +314,17 @@ func resToChatHandleFuncArgsDone(evt *ResponsesStreamEvent, state *ResponsesEven return nil } - if evt.Type == "response.function_call_arguments.done" { - if evt.Arguments == "" { - return nil - } - // Share the exact prefix digest and byte count with output_item.done and - // the final response. A separate accumulated string would duplicate the - // suffix later and keep an unbounded second copy of the arguments. - chunks := resToChatFunctionItem(&ResponsesOutput{ - Type: "function_call", ID: evt.ItemID, CallID: evt.CallID, - Name: evt.Name, Arguments: evt.Arguments, - }, evt.OutputIndex, false, state) - if tool := state.functionCalls[evt.OutputIndex]; tool != nil && state.ProtocolError == "" { - tool.argumentsDone = json.Valid([]byte(evt.Arguments)) - } - return chunks + completed := evt.Arguments + if evt.Type == "response.custom_tool_call_input.done" { + completed = evt.Input } - tool := state.customToolArguments[evt.OutputIndex] - if tool == nil || evt.Input == "" { - return nil - } - remainder, ok := resToChatArgumentsSnapshotSuffix(evt.Input, tool, state) - if !ok { - return nil - } - tool.argumentsDone = true - if remainder == "" { + current := state.OutputIndexToArguments[evt.OutputIndex] + if completed == "" || !strings.HasPrefix(completed, current) || completed == current { return nil } + + remainder := completed[len(current):] + state.OutputIndexToArguments[evt.OutputIndex] = completed return []ChatCompletionsChunk{makeChatDeltaChunk(state, ChatDelta{ ToolCalls: []ChatToolCall{{ Index: &idx, @@ -420,13 +377,7 @@ func resToChatHandleCompleted(evt *ResponsesStreamEvent, state *ResponsesEventTo finishReason = "tool_calls" } - chunks := resToChatCompleteFunctions(evt.Response, state) - if state.ProtocolError != "" { - return nil - } - if evt.Response != nil && evt.Response.Status == "completed" && state.SawToolCall { - finishReason = "tool_calls" - } + var chunks []ChatCompletionsChunk chunks = append(chunks, makeChatFinishChunk(state, finishReason)) if state.IncludeUsage && state.Usage != nil { diff --git a/backend/internal/pkg/openai/request.go b/backend/internal/pkg/openai/request.go index 8ddf8d188d90..15d4bd14cb75 100644 --- a/backend/internal/pkg/openai/request.go +++ b/backend/internal/pkg/openai/request.go @@ -49,7 +49,7 @@ const codexOfficialClientFamilyPrefix = "codex " // 或在 bump context/codex 时同步补入本集合。 var codexOfficialClientOriginators = map[string]bool{ "codex_cli_rs": true, // CLI 默认 DEFAULT_ORIGINATOR - "codex-tui": true, // 交互式 TUI（连字符，真实流量占比最高） + CodexTUIOriginator: true, // 交互式 TUI（连字符，真实流量占比最高） "codex_vscode": true, // VSCode/Cursor 扩展 "codex_vscode_copilot": true, // 扩展 GitHub Copilot 集成模式 "codex_app": true, // 历史保留 @@ -259,11 +259,43 @@ func canonicalizeCodexOriginator(name string) string { return name } -// CodexCLIOriginator 是 codex-rs 客户端的历史默认 originator，保留用于兼容识别。 +// CodexCLIOriginator 是 codex-rs 客户端的默认 originator。 const CodexCLIOriginator = "codex_cli_rs" -// CodexDefaultOriginator 是网关默认使用的 Codex TUI originator。 -const CodexDefaultOriginator = "codex-tui" +// CodexTUIOriginator 是交互式 Codex TUI 使用的 originator。 +const CodexTUIOriginator = "codex-tui" + +// CodexDefaultOriginator 是网关默认发送的 Codex originator。 +// 历史 TUI 身份曾触发上游的身份风控分层，因此默认使用 CLI 身份；TUI 仍保留在官方身份集合中， +// 供入站识别和兼容归一化。 +const CodexDefaultOriginator = CodexCLIOriginator + +// NormalizeCodexClientIdentityToCLI 将已配对的 TUI 身份归一为 CLI 身份。 +// sub2api 历史默认 UA 使用 codex-tui；上游仓库的 #5198 记录了该身份会收到 +// server_is_overloaded，而相同请求改用 codex_cli_rs 可以成功。归一化必须同时改写 +// originator 和 User-Agent 首段，避免触发 #3901 的身份配对校验。 +// +// 版本、操作系统、架构和终端指纹保持不变；尾部 `(name; version)` 客户端标识组 +// 属于 TUI 身份时移除，因为真实 CLI UA 不携带该组。 +func NormalizeCodexClientIdentityToCLI(originator, userAgent string) (string, string, bool) { + if !strings.EqualFold(strings.TrimSpace(originator), CodexTUIOriginator) { + return originator, userAgent, false + } + + ua := strings.TrimSpace(userAgent) + slash := strings.IndexByte(ua, '/') + if slash <= 0 { + return CodexCLIOriginator, ua, true + } + + rest := ua[slash:] + if trailer := codexUATrailerName(ua); trailer != "" && IsCodexOfficialClientOriginator(trailer) { + if open := strings.LastIndex(rest, "("); open > 0 { + rest = strings.TrimRight(rest[:open], " ") + } + } + return CodexCLIOriginator, CodexCLIOriginator + rest, true +} // CodexUserAgentVersion 提取 Codex UA 的完整版本段，即 `{client}/{version} (...` 中的 version。 // 与 ParseCodexEngineVersion 的区别：后者只取三段数字用于引擎版本比较（会丢掉 -alpha.4 diff --git a/backend/internal/pkg/openai/request_codex_identity_normalization_test.go b/backend/internal/pkg/openai/request_codex_identity_normalization_test.go new file mode 100644 index 000000000000..a2d0b357da7f --- /dev/null +++ b/backend/internal/pkg/openai/request_codex_identity_normalization_test.go @@ -0,0 +1,78 @@ +package openai + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestNormalizeCodexClientIdentityToCLI(t *testing.T) { + tests := []struct { + name string + originator string + userAgent string + wantOriginator string + wantUserAgent string + wantChanged bool + }{ + { + name: "TUI 完整身份", + originator: CodexTUIOriginator, + userAgent: "codex-tui/0.163.0 (Windows 11; x86_64) WindowsTerminal (codex-tui; 0.163.0)", + wantOriginator: CodexCLIOriginator, + wantUserAgent: "codex_cli_rs/0.163.0 (Windows 11; x86_64) WindowsTerminal", + wantChanged: true, + }, + { + name: "TUI 只有系统指纹", + originator: CodexTUIOriginator, + userAgent: "codex-tui/0.163.0 (Windows 11; x86_64) WindowsTerminal", + wantOriginator: CodexCLIOriginator, + wantUserAgent: "codex_cli_rs/0.163.0 (Windows 11; x86_64) WindowsTerminal", + wantChanged: true, + }, + { + name: "CLI 身份保持不变", + originator: CodexCLIOriginator, + userAgent: "codex_cli_rs/0.163.0 (Windows 11; x86_64) WindowsTerminal", + wantOriginator: CodexCLIOriginator, + wantUserAgent: "codex_cli_rs/0.163.0 (Windows 11; x86_64) WindowsTerminal", + wantChanged: false, + }, + { + name: "其他官方客户端保持不变", + originator: "codex_vscode", + userAgent: "codex_vscode/0.163.0 (Windows 11; x86_64) vscode (codex_vscode; 0.163.0)", + wantOriginator: "codex_vscode", + wantUserAgent: "codex_vscode/0.163.0 (Windows 11; x86_64) vscode (codex_vscode; 0.163.0)", + wantChanged: false, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + originator, userAgent, changed := NormalizeCodexClientIdentityToCLI(test.originator, test.userAgent) + require.Equal(t, test.wantOriginator, originator) + require.Equal(t, test.wantUserAgent, userAgent) + require.Equal(t, test.wantChanged, changed) + }) + } +} + +func TestNormalizeCodexClientIdentityToCLIIsPairedAndIdempotent(t *testing.T) { + originator, userAgent, changed := NormalizeCodexClientIdentityToCLI( + CodexTUIOriginator, + "codex-tui/0.163.0 (Windows 11; x86_64) WindowsTerminal (codex-tui; 0.163.0)", + ) + require.True(t, changed) + + pairedOriginator, pairedUserAgent, ok := PairCodexClientIdentity(userAgent) + require.True(t, ok) + require.Equal(t, originator, pairedOriginator) + require.Equal(t, userAgent, pairedUserAgent) + + againOriginator, againUserAgent, changedAgain := NormalizeCodexClientIdentityToCLI(originator, userAgent) + require.False(t, changedAgain) + require.Equal(t, originator, againOriginator) + require.Equal(t, userAgent, againUserAgent) +} diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index 808bd1a5088d..65b275671c4e 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -5,30 +5,39 @@ import ( "crypto/rand" "encoding/hex" "encoding/json" + "errors" "fmt" "log/slog" + "slices" "strconv" + "sync" + "sync/atomic" "time" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/redis/go-redis/v9" + "github.com/tidwall/gjson" ) const ( - schedulerBucketSetKey = "sched:buckets" - schedulerOutboxWatermarkKey = "sched:outbox:watermark" - schedulerAccountPrefix = "sched:acc:" - schedulerAccountMetaPrefix = "sched:meta:v2:" - schedulerLegacyAccountMetaPrefix = "sched:meta:" - schedulerAccountLastUsedPrefix = "sched:acc:last_used:" - schedulerActivePrefix = "sched:active:v2:" - schedulerReadyPrefix = "sched:ready:v2:" - schedulerLegacyReadyPrefix = "sched:ready:" - schedulerVersionPrefix = "sched:ver:v2:" - schedulerEpochPrefix = "sched:epoch:" - schedulerRetiredPrefix = "sched:retired:" - schedulerSnapshotPrefix = "sched:v2:" - schedulerLockPrefix = "sched:lock:" + schedulerBucketSetKey = "sched:buckets" + schedulerOutboxWatermarkKey = "sched:outbox:watermark" + schedulerAccountPrefix = "sched:acc:" + schedulerAccountMetaPrefix = "sched:meta:" + schedulerAccountLastUsedPrefix = "sched:acc:last_used:" + schedulerActivePrefix = "sched:active:" + schedulerReadyPrefix = "sched:ready:" + schedulerVersionPrefix = "sched:ver:" + schedulerEpochPrefix = "sched:epoch:" + schedulerRetiredPrefix = "sched:retired:" + schedulerSnapshotPrefix = "sched:" + schedulerLockPrefix = "sched:lock:" + // 账号精简记录（sched:meta:）的修订号：全局修订号与按平台划分的修订号（前缀后接平台名）。 + // 改写精简记录后写入所涉平台的修订号；删除记录、改写平台为空或平台发生变化的记录、写入出错时写入全局修订号。 + // 取值都是新的随机串，读侧只比较是否相等，键丢失后重建也不会与旧取值相同； + // Redis 回滚恢复出缓存记录过的旧取值时，由最长复用时长兜底。 + schedulerAccountMetaRevisionKey = "sched:meta:rev" + schedulerAccountMetaPlatformRevisionPrefix = "sched:meta:rev:" defaultSchedulerSnapshotMGetChunkSize = 128 defaultSchedulerSnapshotWriteChunkSize = 256 @@ -37,6 +46,16 @@ const ( // snapshotGraceTTLSeconds 旧快照过期的宽限期（秒）。 // 替代立即 DEL，让正在读取旧版本的 reader 有足够时间完成 ZRANGE。 snapshotGraceTTLSeconds = 60 + + // schedulerSnapshotCacheMaxAge 是进程内快照缓存条目的最长复用时长， + // 限定了未推进修订号的精简记录改写在本进程的最长可见延迟，与 outbox_lag_rebuild_seconds 默认 10 秒的陈旧容忍度相同。 + schedulerSnapshotCacheMaxAge = 10 * time.Second + // schedulerSnapshotCacheEvictAge：写入条目时，加载时间早于该时长的其它条目一并清除， + // 已退休分组的条目不会长期驻留。 + schedulerSnapshotCacheEvictAge = 10 * schedulerSnapshotCacheMaxAge + + schedulerAccountMetaRevisionBytes = 16 + schedulerAccountMetaRevisionWriteTimeout = 3 * time.Second ) const ( @@ -121,7 +140,7 @@ local currentActive = redis.call('GET', KEYS[5]) if currentActive ~= false then redis.call('EXPIRE', ARGV[2] .. currentActive, tonumber(ARGV[3])) end -redis.call('DEL', KEYS[4], KEYS[5], KEYS[6]) +redis.call('DEL', KEYS[4], KEYS[5]) return currentEpoch `) @@ -129,7 +148,6 @@ return currentEpoch local currentEpochRaw = redis.call('GET', KEYS[1]) local currentEpoch = tonumber(currentEpochRaw) local retiredEpochRaw = redis.call('GET', KEYS[2]) -redis.call('DEL', KEYS[6]) if retiredEpochRaw == false then if currentEpochRaw == false then @@ -173,13 +191,12 @@ return 0 // 仅当新版本号 >= 当前激活版本时才切换，防止并发写入导致版本回滚。 // 旧快照使用 EXPIRE 设置宽限期而非立即 DEL，避免与 reader 竞态。 // - // KEYS[1] = activeKey (sched:active:v2:{bucket}) - // KEYS[2] = readyKey (sched:ready:v2:{bucket}) + // KEYS[1] = activeKey (sched:active:{bucket}) + // KEYS[2] = readyKey (sched:ready:{bucket}) // KEYS[3] = bucketSetKey (sched:buckets) // KEYS[4] = snapshotKey (新写入的快照 key) // KEYS[5] = epochKey // KEYS[6] = retiredKey - // KEYS[7] = legacyReadyKey (sched:ready:{bucket}) // ARGV[1] = 新版本号字符串 // ARGV[2] = bucket 字符串 (用于 SADD) // ARGV[3] = 快照 key 前缀 (用于构造旧快照 key) @@ -214,7 +231,6 @@ end redis.call('SET', KEYS[1], ARGV[1]) redis.call('SET', KEYS[2], '1') redis.call('SADD', KEYS[3], ARGV[2]) -redis.call('DEL', KEYS[7]) if currentActive ~= false and currentActive ~= ARGV[1] then redis.call('EXPIRE', ARGV[3] .. currentActive, tonumber(ARGV[4])) @@ -228,6 +244,28 @@ type schedulerCache struct { rdb *redis.Client mgetChunkSize int writeChunkSize int + + // now 是进程内快照缓存判定条目年龄的时钟。 + now func() time.Time + + snapshotMu sync.RWMutex + snapshotEntries map[string]*schedulerSnapshotEntry + + // globalRevisionPending 与 globalRevisionWriting 合并写入出错后的后台全局修订号写入，见 scheduleGlobalAccountMetaRevision。 + globalRevisionPending atomic.Bool + globalRevisionWriting atomic.Bool +} + +// schedulerSnapshotEntry 是一个 bucket 的已解码快照。 +// revisions 依次是全局修订号与该 bucket 成员可能涉及的各平台修订号（见 schedulerSnapshotTagKeys）。 +// accounts 是按快照顺序排列的母本，发布后只读，指针不外泄； +// 每次读取都按值复制出本次请求独占的副本再套用 last_used。 +type schedulerSnapshotEntry struct { + active string + revisions []string + loadedAt time.Time + accounts []*service.Account + lastUsedKeys []string } func NewSchedulerCache(rdb *redis.Client) service.SchedulerCache { @@ -242,42 +280,67 @@ func newSchedulerCacheWithChunkSizes(rdb *redis.Client, mgetChunkSize, writeChun writeChunkSize = defaultSchedulerSnapshotWriteChunkSize } return &schedulerCache{ - rdb: rdb, - mgetChunkSize: mgetChunkSize, - writeChunkSize: writeChunkSize, + rdb: rdb, + mgetChunkSize: mgetChunkSize, + writeChunkSize: writeChunkSize, + now: time.Now, + snapshotEntries: make(map[string]*schedulerSnapshotEntry), } } +// GetSnapshot 读取 bucket 的激活快照。解码后的精简记录按 bucket 常驻进程内存， +// 只有 ready、active 版本、全局修订号与该 bucket 所涉平台的修订号都与 Redis 一致且未超过最长复用时长时才复用； +// 否则重新从 Redis 加载。last_used 每次都从 Redis 读取。 +// 返回的账号对象由本次调用独占，可以修改其顶层字段；map、切片底层数组与指针指向的值与其它调用共享，不得原地修改。 func (c *schedulerCache) GetSnapshot(ctx context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) { - readyKey := schedulerBucketKey(schedulerReadyPrefix, bucket) - readyVal, err := c.rdb.Get(ctx, readyKey).Result() - if err == redis.Nil { - return nil, false, nil - } + bucketKey := bucket.String() + tagKeys := schedulerSnapshotTagKeys(bucket) + tags, err := c.rdb.MGet(ctx, tagKeys...).Result() if err != nil { return nil, false, err } - if readyVal != "1" { - return nil, false, nil + if len(tags) != len(tagKeys) { + return nil, false, fmt.Errorf("unexpected scheduler snapshot tag count: got %d, want %d", len(tags), len(tagKeys)) } - - activeKey := schedulerBucketKey(schedulerActivePrefix, bucket) - activeVal, err := c.rdb.Get(ctx, activeKey).Result() - if err == redis.Nil { - return nil, false, nil + readyVal, err := schedulerCacheOptionalString(tags[0]) + if err != nil { + return nil, false, err } + activeVal, err := schedulerCacheOptionalString(tags[1]) if err != nil { return nil, false, err } + // 修订号缺失（nil）按空字符串比较：它是一个独立取值，不会与任何写入过的随机修订号相等。 + revisions := make([]string, 0, len(tags)-2) + for _, tag := range tags[2:] { + revision, err := schedulerCacheOptionalString(tag) + if err != nil { + return nil, false, err + } + revisions = append(revisions, revision) + } + if readyVal != "1" || activeVal == "" { + c.dropSnapshotEntry(bucketKey) + return nil, false, nil + } - snapshotKey := schedulerSnapshotKey(bucket, activeVal) - ids, err := c.rdb.ZRange(ctx, snapshotKey, 0, -1).Result() + now := c.now() + if entry := c.lookupSnapshotEntry(bucketKey); entry != nil && entry.matches(activeVal, revisions, now) { + accounts, err := c.snapshotAccountsFromEntry(ctx, entry) + if err != nil { + return nil, false, err + } + return accounts, true, nil + } + + ids, err := c.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, activeVal), 0, -1).Result() if err != nil { return nil, false, err } if len(ids) == 0 { // 空快照视为缓存未命中，触发数据库回退查询 // 这解决了新分组创建后立即绑定账号时的竞态条件问题 + c.dropSnapshotEntry(bucketKey) return nil, false, nil } @@ -291,27 +354,141 @@ func (c *schedulerCache) GetSnapshot(ctx context.Context, bucket service.Schedul if err != nil { return nil, false, err } - lastUsedValues, err := c.mgetChunked(ctx, lastUsedKeys) - if err != nil { - return nil, false, err - } - accounts := make([]*service.Account, 0, len(values)) - for i, val := range values { + masters := make([]*service.Account, 0, len(values)) + for _, val := range values { if val == nil { + c.dropSnapshotEntry(bucketKey) return nil, false, nil } account, err := decodeCachedAccount(val) if err != nil { return nil, false, err } + masters = append(masters, account) + } + // 母本发布进缓存前由本 goroutine 完成预处理，发布后只读。 + for _, master := range masters { + master.PrepareSharedSnapshot() + } + + // 条目记录的是加载前读到的 (active, 修订号) 与当时的时钟：精简记录读取晚于这组取值， + // 之后若再有改写，修订号或 active 必然不同，或在最长复用时长后过期。 + entry := &schedulerSnapshotEntry{ + active: activeVal, + revisions: revisions, + loadedAt: now, + accounts: masters, + lastUsedKeys: lastUsedKeys, + } + c.storeSnapshotEntry(bucketKey, entry) + + accounts, err := c.snapshotAccountsFromEntry(ctx, entry) + if err != nil { + return nil, false, err + } + return accounts, true, nil +} + +func (e *schedulerSnapshotEntry) matches(active string, revisions []string, now time.Time) bool { + if e.active != active || !slices.Equal(e.revisions, revisions) { + return false + } + age := now.Sub(e.loadedAt) + return age >= 0 && age < schedulerSnapshotCacheMaxAge +} + +// snapshotAccountsFromEntry 读取最新 last_used，并为本次调用一次性分配一组母本副本。 +func (c *schedulerCache) snapshotAccountsFromEntry(ctx context.Context, entry *schedulerSnapshotEntry) ([]*service.Account, error) { + lastUsedValues, err := c.mgetChunked(ctx, entry.lastUsedKeys) + if err != nil { + return nil, err + } + if len(lastUsedValues) != len(entry.accounts) { + return nil, fmt.Errorf("unexpected scheduler last_used value count: got %d, want %d", len(lastUsedValues), len(entry.accounts)) + } + slab := make([]service.Account, len(entry.accounts)) + accounts := make([]*service.Account, len(entry.accounts)) + for i, master := range entry.accounts { + slab[i] = *master + account := &slab[i] if err := applySchedulerLastUsed(account, lastUsedValues[i]); err != nil { - return nil, false, err + return nil, err } - accounts = append(accounts, account) + accounts[i] = account } + return accounts, nil +} - return accounts, true, nil +func (c *schedulerCache) lookupSnapshotEntry(bucketKey string) *schedulerSnapshotEntry { + c.snapshotMu.RLock() + defer c.snapshotMu.RUnlock() + return c.snapshotEntries[bucketKey] +} + +// storeSnapshotEntry 写入条目，并清除加载时间早于 schedulerSnapshotCacheEvictAge 的其它条目。 +// 已有条目的加载时间晚于新条目时保留已有条目：它基于更晚读到的 (active, 修订号)。 +func (c *schedulerCache) storeSnapshotEntry(bucketKey string, entry *schedulerSnapshotEntry) { + c.snapshotMu.Lock() + defer c.snapshotMu.Unlock() + if current := c.snapshotEntries[bucketKey]; current == nil || !current.loadedAt.After(entry.loadedAt) { + c.snapshotEntries[bucketKey] = entry + } + for key, other := range c.snapshotEntries { + if key != bucketKey && entry.loadedAt.Sub(other.loadedAt) > schedulerSnapshotCacheEvictAge { + delete(c.snapshotEntries, key) + } + } +} + +func (c *schedulerCache) dropSnapshotEntry(bucketKey string) { + c.snapshotMu.RLock() + _, ok := c.snapshotEntries[bucketKey] + c.snapshotMu.RUnlock() + if !ok { + return + } + c.snapshotMu.Lock() + delete(c.snapshotEntries, bucketKey) + c.snapshotMu.Unlock() +} + +// schedulerSnapshotTagKeys 返回 GetSnapshot 一次读取的键：ready、active、全局修订号、bucket 平台的修订号， +// mixed 模式再加 antigravity 平台的修订号。按平台判定失效依赖 bucket 成员只取自 bucket 平台 +// （mixed 模式另含开启混合调度的 antigravity 账号）；账号平台发生变化时，写入方另行写入全局修订号。 +func schedulerSnapshotTagKeys(bucket service.SchedulerBucket) []string { + keys := make([]string, 0, 5) + keys = append(keys, + schedulerBucketKey(schedulerReadyPrefix, bucket), + schedulerBucketKey(schedulerActivePrefix, bucket), + schedulerAccountMetaRevisionKey, + schedulerAccountMetaRevisionKeyForPlatform(bucket.Platform), + ) + if bucket.Mode == service.SchedulerModeMixed && bucket.Platform != service.PlatformAntigravity { + keys = append(keys, schedulerAccountMetaRevisionKeyForPlatform(service.PlatformAntigravity)) + } + return keys +} + +// schedulerAccountMetaRevisionKeyForPlatform 返回平台修订号的键；平台为空时返回全局修订号的键。 +func schedulerAccountMetaRevisionKeyForPlatform(platform string) string { + if platform == "" { + return schedulerAccountMetaRevisionKey + } + return schedulerAccountMetaPlatformRevisionPrefix + platform +} + +func schedulerCacheOptionalString(value any) (string, error) { + switch typed := value.(type) { + case nil: + return "", nil + case string: + return typed, nil + case []byte: + return string(typed), nil + default: + return "", fmt.Errorf("unexpected scheduler cache value type: %T", value) + } } func (c *schedulerCache) CaptureBucketWriteToken(ctx context.Context, bucket service.SchedulerBucket) (service.SchedulerBucketWriteToken, error) { @@ -329,14 +506,13 @@ func (c *schedulerCache) CaptureBucketWriteToken(ctx context.Context, bucket ser } func (c *schedulerCache) RetireBucket(ctx context.Context, bucket service.SchedulerBucket) error { - snapshotKeyPrefix := schedulerSnapshotKeyPrefix(bucket) + snapshotKeyPrefix := fmt.Sprintf("%s%d:%s:%s:v", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode) result, err := retireBucketScript.Run(ctx, c.rdb, []string{ schedulerBucketKey(schedulerEpochPrefix, bucket), schedulerBucketKey(schedulerRetiredPrefix, bucket), schedulerBucketSetKey, schedulerBucketKey(schedulerReadyPrefix, bucket), schedulerBucketKey(schedulerActivePrefix, bucket), - schedulerBucketKey(schedulerLegacyReadyPrefix, bucket), }, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds).Int64() if err != nil { return err @@ -348,14 +524,13 @@ func (c *schedulerCache) RetireBucket(ctx context.Context, bucket service.Schedu } func (c *schedulerCache) ReopenBucket(ctx context.Context, bucket service.SchedulerBucket) (service.SchedulerBucketWriteToken, error) { - snapshotKeyPrefix := schedulerSnapshotKeyPrefix(bucket) + snapshotKeyPrefix := fmt.Sprintf("%s%d:%s:%s:v", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode) result, err := reopenBucketScript.Run(ctx, c.rdb, []string{ schedulerBucketKey(schedulerEpochPrefix, bucket), schedulerBucketKey(schedulerRetiredPrefix, bucket), schedulerBucketSetKey, schedulerBucketKey(schedulerReadyPrefix, bucket), schedulerBucketKey(schedulerActivePrefix, bucket), - schedulerBucketKey(schedulerLegacyReadyPrefix, bucket), }, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds).Int64() if err != nil { return service.SchedulerBucketWriteToken{}, err @@ -543,7 +718,7 @@ func (c *schedulerCache) activateSnapshotVersion(ctx context.Context, bucket ser // 旧快照使用 EXPIRE 宽限期而非立即 DEL，避免 reader 竞态。 activeKey := schedulerBucketKey(schedulerActivePrefix, bucket) readyKey := schedulerBucketKey(schedulerReadyPrefix, bucket) - snapshotKeyPrefix := schedulerSnapshotKeyPrefix(bucket) + snapshotKeyPrefix := fmt.Sprintf("%s%d:%s:%s:v", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode) keys := []string{ activeKey, @@ -552,7 +727,6 @@ func (c *schedulerCache) activateSnapshotVersion(ctx context.Context, bucket ser snapshotKey, schedulerBucketKey(schedulerEpochPrefix, bucket), schedulerBucketKey(schedulerRetiredPrefix, bucket), - schedulerBucketKey(schedulerLegacyReadyPrefix, bucket), } args := []any{version, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds, token.Epoch} @@ -612,13 +786,13 @@ func (c *schedulerCache) DeleteAccount(ctx context.Context, accountID int64) err return nil } id := strconv.FormatInt(accountID, 10) - return c.rdb.Del( - ctx, - schedulerAccountKey(id), - schedulerAccountMetaKey(id), - schedulerLegacyAccountMetaKey(id), - schedulerLastUsedKey(id), - ).Err() + // 删除精简记录与推进全局修订号在同一个 MULTI 中原子执行：读到删除结果的读者也读到新修订号， + // 进程内快照缓存不会给出已删除的账号。 + pipe := c.rdb.TxPipeline() + pipe.Del(ctx, schedulerAccountKey(id), schedulerAccountMetaKey(id), schedulerLastUsedKey(id)) + pipe.Set(ctx, schedulerAccountMetaRevisionKey, newSchedulerAccountMetaRevision(), 0) + _, err := pipe.Exec(ctx) + return err } func (c *schedulerCache) UpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error { @@ -628,6 +802,7 @@ func (c *schedulerCache) UpdateLastUsed(ctx context.Context, updates map[int64]t pipe := c.rdb.Pipeline() queued := 0 + metaDeleted := false keys := make([]string, 0, schedulerLastUsedUpdateChunkSize*2) args := make([]any, 0, schedulerLastUsedUpdateChunkSize) queueBatch := func() { @@ -650,14 +825,9 @@ func (c *schedulerCache) UpdateLastUsed(ctx context.Context, updates map[int64]t "error", err, ) idText := strconv.FormatInt(id, 10) - pipe.Del( - ctx, - schedulerAccountKey(idText), - schedulerAccountMetaKey(idText), - schedulerLegacyAccountMetaKey(idText), - schedulerLastUsedKey(idText), - ) + pipe.Del(ctx, schedulerAccountKey(idText), schedulerAccountMetaKey(idText), schedulerLastUsedKey(idText)) queued++ + metaDeleted = true continue } idText := strconv.FormatInt(id, 10) @@ -671,6 +841,10 @@ func (c *schedulerCache) UpdateLastUsed(ctx context.Context, updates map[int64]t if queued == 0 { return nil } + if metaDeleted { + // 被删账号的平台未知，全局修订号排在同一 pipeline 的 DEL 之后。 + pipe.Set(ctx, schedulerAccountMetaRevisionKey, newSchedulerAccountMetaRevision(), 0) + } _, err := pipe.Exec(ctx) return err } @@ -721,7 +895,7 @@ func (c *schedulerCache) SetOutboxWatermark(ctx context.Context, id int64) error } func schedulerBucketKey(prefix string, bucket service.SchedulerBucket) string { - return prefix + bucket.String() + return fmt.Sprintf("%s%d:%s:%s", prefix, bucket.GroupID, bucket.Platform, bucket.Mode) } func schedulerGroupLifecycleLockKey(groupID int64) string { @@ -729,11 +903,7 @@ func schedulerGroupLifecycleLockKey(groupID int64) string { } func schedulerSnapshotKey(bucket service.SchedulerBucket, version string) string { - return schedulerSnapshotKeyPrefix(bucket) + version -} - -func schedulerSnapshotKeyPrefix(bucket service.SchedulerBucket) string { - return schedulerSnapshotPrefix + bucket.String() + ":v" + return fmt.Sprintf("%s%d:%s:%s:v%s", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode, version) } func schedulerAccountKey(id string) string { @@ -744,10 +914,6 @@ func schedulerAccountMetaKey(id string) string { return schedulerAccountMetaPrefix + id } -func schedulerLegacyAccountMetaKey(id string) string { - return schedulerLegacyAccountMetaPrefix + id -} - func schedulerLastUsedKey(id string) string { return schedulerAccountLastUsedPrefix + id } @@ -804,25 +970,54 @@ func decodeCachedAccount(val any) (*service.Account, error) { return &account, nil } +// writeAccountIDs 分块写入整账号与精简记录，返回成功编码并写入的有序账号 ID。 +// 每块在追加下一个账号前才执行，最后一块总是非空；所涉平台的修订号追加在最后一块末尾， +// 与最后一批精简记录同一次往返发出，且排在全部精简记录写入之后。 +// 精简记录用 SET ... GET 写入并取回旧值：旧值的平台与新平台不同时，旧平台 bucket 中可能仍缓存着该账号， +// 全部写入完成后另行写入全局修订号。该检测尽力而为（客户端在 pipeline 内部重试时取回的可能已是本次写入的值）， +// 漏检由进程内快照缓存的最长复用时长兜底。 +// 任一块执行出错时，此前已执行的块与出错的块都可能已有精简记录生效，返回时启动后台写入全局修订号。 func (c *schedulerCache) writeAccountIDs(ctx context.Context, accounts []service.Account) ([]int64, error) { if len(accounts) == 0 { return nil, nil } - // Each chunk must atomically publish the new full/meta payloads and remove - // the legacy metadata copy. A regular pipeline may partially execute across - // a connection failure, leaving a retired credential readable after rollback. - pipe := c.rdb.TxPipeline() + type metaWrite struct { + platform string + previous *redis.StatusCmd + } + pipe := c.rdb.Pipeline() accountIDs := make([]int64, 0, len(accounts)) + platforms := make([]string, 0, 1) + writes := make([]metaWrite, 0, min(len(accounts), c.writeChunkSize)) pending := 0 - flush := func() error { - if pending == 0 { - return nil + execFailed := false + platformChanged := false + defer func() { + if execFailed { + c.scheduleGlobalAccountMetaRevision(ctx) } - if _, err := pipe.Exec(ctx); err != nil { + }() + exec := func() error { + cmds, err := pipe.Exec(ctx) + if err != nil && !errors.Is(err, redis.Nil) { + execFailed = true return err } - pipe = c.rdb.TxPipeline() + // SET ... GET 在键原本不存在时返回 redis.Nil，不算失败；其余命令的错误逐条检查。 + for _, cmd := range cmds { + if cmdErr := cmd.Err(); cmdErr != nil && !errors.Is(cmdErr, redis.Nil) { + execFailed = true + return cmdErr + } + } + for _, write := range writes { + if schedulerMetaPlatformChanged(write.previous, write.platform) { + platformChanged = true + } + } + writes = writes[:0] + pipe = c.rdb.Pipeline() pending = 0 return nil } @@ -834,49 +1029,95 @@ func (c *schedulerCache) writeAccountIDs(ctx context.Context, accounts []service "account_id", account.ID, "error", err, ) - if account.ID > 0 { - id := strconv.FormatInt(account.ID, 10) - pipe.Del( - ctx, - schedulerAccountKey(id), - schedulerAccountMetaKey(id), - schedulerLegacyAccountMetaKey(id), - schedulerLastUsedKey(id), - ) - pending++ - if pending >= c.writeChunkSize { - if err := flush(); err != nil { - return nil, err - } - } - } continue } + if pending >= c.writeChunkSize { + if err := exec(); err != nil { + return nil, err + } + } id := strconv.FormatInt(account.ID, 10) pipe.Set(ctx, schedulerAccountKey(id), fullPayload, 0) - pipe.Set(ctx, schedulerAccountMetaKey(id), metaPayload, 0) - // v2 metadata is generation-isolated. Remove the prior generation in the - // same pipeline so credential rotation and rebuilds cannot retain an old - // API key indefinitely under sched:meta:. - pipe.Del(ctx, schedulerLegacyAccountMetaKey(id)) + writes = append(writes, metaWrite{ + platform: account.Platform, + previous: pipe.SetArgs(ctx, schedulerAccountMetaKey(id), metaPayload, redis.SetArgs{Get: true}), + }) // Keep the hot LastUsedAt side key untouched: a lagging snapshot rebuild // must not overwrite a newer scheduler update. accountIDs = append(accountIDs, account.ID) - pending++ - if pending >= c.writeChunkSize { - if err := flush(); err != nil { - return nil, err - } + if !slices.Contains(platforms, account.Platform) { + platforms = append(platforms, account.Platform) } + pending++ } - if err := flush(); err != nil { + if pending == 0 { + return accountIDs, nil + } + for _, platform := range platforms { + pipe.Set(ctx, schedulerAccountMetaRevisionKeyForPlatform(platform), newSchedulerAccountMetaRevision(), 0) + } + if err := exec(); err != nil { return nil, err } + if platformChanged { + c.advanceGlobalAccountMetaRevision(ctx) + } return accountIDs, nil } +// schedulerMetaPlatformChanged 判断 SET ... GET 取回的旧精简记录是否属于另一个平台； +// 旧值不是合法 JSON 或缺少字符串类型的顶层 Platform 时按已变化处理。 +func schedulerMetaPlatformChanged(previous *redis.StatusCmd, platform string) bool { + raw, err := previous.Result() + if err != nil { + return false + } + if !gjson.Valid(raw) { + return true + } + old := gjson.Get(raw, "Platform") + return old.Type != gjson.String || old.Str != platform +} + +// scheduleGlobalAccountMetaRevision 在后台写入全局修订号，不阻塞调用方。 +// 同一时刻至多一个后台写入：写入前先清除待写标记，写完后标记再次被置位就再写一轮， +// 因此每次调用之后都至少还有一次完整的写入发生在它之后。 +func (c *schedulerCache) scheduleGlobalAccountMetaRevision(ctx context.Context) { + c.globalRevisionPending.Store(true) + if !c.globalRevisionWriting.CompareAndSwap(false, true) { + return + } + go func() { + for { + c.globalRevisionPending.Store(false) + c.advanceGlobalAccountMetaRevision(ctx) + c.globalRevisionWriting.Store(false) + if !c.globalRevisionPending.Load() || !c.globalRevisionWriting.CompareAndSwap(false, true) { + return + } + } + }() +} + +// advanceGlobalAccountMetaRevision 单独写入新的全局修订号，使所有 bucket 的进程内快照失效。 +// 不受调用方取消影响；失败只记录日志，由进程内快照缓存的最长复用时长兜底。 +func (c *schedulerCache) advanceGlobalAccountMetaRevision(ctx context.Context) { + writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), schedulerAccountMetaRevisionWriteTimeout) + defer cancel() + if err := c.rdb.Set(writeCtx, schedulerAccountMetaRevisionKey, newSchedulerAccountMetaRevision(), 0).Err(); err != nil { + slog.Warn("scheduler cache failed to advance account metadata revision", "error", err) + } +} + +func newSchedulerAccountMetaRevision() string { + raw := make([]byte, schedulerAccountMetaRevisionBytes) + // crypto/rand.Read 读取失败时直接终止进程，不会返回错误。 + _, _ = rand.Read(raw) + return hex.EncodeToString(raw) +} + func marshalSchedulerCacheAccount(account service.Account) ([]byte, []byte, error) { fullPayload, err := json.Marshal(account) if err != nil { diff --git a/backend/internal/repository/scheduler_cache_snapshot_memo_unit_test.go b/backend/internal/repository/scheduler_cache_snapshot_memo_unit_test.go new file mode 100644 index 000000000000..9b3d92ea2635 --- /dev/null +++ b/backend/internal/repository/scheduler_cache_snapshot_memo_unit_test.go @@ -0,0 +1,980 @@ +//go:build unit + +package repository + +import ( + "context" + "errors" + "fmt" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +// schedulerCacheCommandRecorder 记录发往 Redis 的命令（名称 + 首个参数）与每次 pipeline 执行的命令批次， +// 并可让第 N 次 pipeline 执行失败。新连接的握手命令不计入。 +type schedulerCacheCommandRecorder struct { + mu sync.Mutex + commands []string + batches [][]string + pipelines int + failPipelineAt int +} + +func (r *schedulerCacheCommandRecorder) DialHook(next redis.DialHook) redis.DialHook { return next } + +func (r *schedulerCacheCommandRecorder) ProcessHook(next redis.ProcessHook) redis.ProcessHook { + return func(ctx context.Context, cmd redis.Cmder) error { + if !isSchedulerCacheHandshakeCommand(cmd) { + r.record(cmd) + } + return next(ctx, cmd) + } +} + +func (r *schedulerCacheCommandRecorder) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook { + return func(ctx context.Context, cmds []redis.Cmder) error { + handshake := len(cmds) > 0 + for _, cmd := range cmds { + handshake = handshake && isSchedulerCacheHandshakeCommand(cmd) + } + if handshake { + return next(ctx, cmds) + } + r.mu.Lock() + r.pipelines++ + fail := r.failPipelineAt > 0 && r.pipelines == r.failPipelineAt + r.mu.Unlock() + batch := make([]string, 0, len(cmds)) + for _, cmd := range cmds { + batch = append(batch, r.record(cmd)) + } + r.mu.Lock() + r.batches = append(r.batches, batch) + r.mu.Unlock() + if fail { + err := errors.New("injected pipeline failure") + for _, cmd := range cmds { + cmd.SetErr(err) + } + return err + } + return next(ctx, cmds) + } +} + +func isSchedulerCacheHandshakeCommand(cmd redis.Cmder) bool { + switch strings.ToLower(cmd.Name()) { + case "hello", "auth", "client", "select": + return true + default: + return false + } +} + +func (r *schedulerCacheCommandRecorder) record(cmd redis.Cmder) string { + entry := strings.ToLower(cmd.Name()) + if args := cmd.Args(); len(args) > 1 { + entry += " " + fmt.Sprint(args[1]) + } + r.mu.Lock() + r.commands = append(r.commands, entry) + r.mu.Unlock() + return entry +} + +func (r *schedulerCacheCommandRecorder) reset() { + r.mu.Lock() + r.commands = nil + r.batches = nil + r.pipelines = 0 + r.mu.Unlock() +} + +func (r *schedulerCacheCommandRecorder) pipelineBatches() [][]string { + r.mu.Lock() + defer r.mu.Unlock() + return append([][]string(nil), r.batches...) +} + +func (r *schedulerCacheCommandRecorder) setFailPipelineAt(n int) { + r.mu.Lock() + r.failPipelineAt = n + r.mu.Unlock() +} + +func (r *schedulerCacheCommandRecorder) recorded() []string { + r.mu.Lock() + defer r.mu.Unlock() + return append([]string(nil), r.commands...) +} + +func (r *schedulerCacheCommandRecorder) countPrefix(prefix string) int { + count := 0 + for _, command := range r.recorded() { + if strings.HasPrefix(command, prefix) { + count++ + } + } + return count +} + +func newSchedulerCacheMemoUnit(t *testing.T, writeChunkSize int) (*schedulerCache, *miniredis.Miniredis, *schedulerCacheCommandRecorder) { + t.Helper() + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + recorder := &schedulerCacheCommandRecorder{} + rdb.AddHook(recorder) + cache, ok := newSchedulerCacheWithChunkSizes(rdb, defaultSchedulerSnapshotMGetChunkSize, writeChunkSize).(*schedulerCache) + require.True(t, ok) + return cache, mr, recorder +} + +func memoTestBucket(groupID int64) service.SchedulerBucket { + return service.SchedulerBucket{GroupID: groupID, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeSingle} +} + +func memoTestAccount(id int64, name string) service.Account { + return memoTestAccountOn(service.PlatformAnthropic, id, name) +} + +func memoTestAccountOn(platform string, id int64, name string) service.Account { + return service.Account{ + ID: id, + Name: name, + Platform: platform, + Type: service.AccountTypeOAuth, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 2, + Credentials: map[string]any{ + "model_mapping": map[string]any{"claude-sonnet": "sonnet-" + name}, + "access_token": "dropped-by-projection", + }, + Extra: map[string]any{"session_window_utilization": 0.25}, + } +} + +func memoPublish(t *testing.T, cache *schedulerCache, bucket service.SchedulerBucket, accounts ...service.Account) { + t.Helper() + ctx := context.Background() + token, err := cache.CaptureBucketWriteToken(ctx, bucket) + require.NoError(t, err) + require.NoError(t, cache.SetSnapshot(ctx, bucket, token, accounts)) +} + +func memoRead(t *testing.T, cache *schedulerCache, bucket service.SchedulerBucket) []*service.Account { + t.Helper() + accounts, hit, err := cache.GetSnapshot(context.Background(), bucket) + require.NoError(t, err) + require.True(t, hit) + return accounts +} + +func memoRequireMiss(t *testing.T, cache *schedulerCache, bucket service.SchedulerBucket) { + t.Helper() + accounts, hit, err := cache.GetSnapshot(context.Background(), bucket) + require.NoError(t, err) + require.False(t, hit) + require.Nil(t, accounts) +} + +func memoIDs(accounts []*service.Account) []int64 { + ids := make([]int64, 0, len(accounts)) + for _, account := range accounts { + ids = append(ids, account.ID) + } + return ids +} + +func memoRevision(t *testing.T, cache *schedulerCache, key string) string { + t.Helper() + revision, err := cache.rdb.Get(context.Background(), key).Result() + if errors.Is(err, redis.Nil) { + return "" + } + require.NoError(t, err) + return revision +} + +// memoRevisionValue 读取修订号，不做断言，供 require.Eventually 的条件函数使用。 +func memoRevisionValue(cache *schedulerCache, key string) string { + revision, _ := cache.rdb.Get(context.Background(), key).Result() + return revision +} + +var memoAnthropicRevisionKey = schedulerAccountMetaRevisionKeyForPlatform(service.PlatformAnthropic) + +func memoEntryCount(cache *schedulerCache) int { + cache.snapshotMu.RLock() + defer cache.snapshotMu.RUnlock() + return len(cache.snapshotEntries) +} + +func TestSchedulerCacheSnapshotMemoSecondReadOnlyFetchesTagsAndLastUsed(t *testing.T) { + cache, _, recorder := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(101) + memoPublish(t, cache, bucket, memoTestAccount(1, "a"), memoTestAccount(2, "b"), memoTestAccount(3, "c")) + + recorder.reset() + first := memoRead(t, cache, bucket) + require.Equal(t, 1, recorder.countPrefix("zrange ")) + require.Equal(t, 1, recorder.countPrefix("mget "+schedulerAccountMetaPrefix)) + + recorder.reset() + second := memoRead(t, cache, bucket) + require.Equal(t, []string{ + "mget " + schedulerBucketKey(schedulerReadyPrefix, bucket), + "mget " + schedulerLastUsedKey("1"), + }, recorder.recorded()) + + require.Equal(t, []int64{1, 2, 3}, memoIDs(second)) + require.Equal(t, first, second) + for i := range first { + require.NotSame(t, first[i], second[i]) + } + require.Equal(t, "sonnet-b", second[1].GetMappedModel("claude-sonnet")) + require.NotContains(t, second[1].Credentials, "access_token") +} + +func TestSchedulerCacheSnapshotMemoLastUsedIsFreshAndCopiesAreIndependent(t *testing.T) { + ctx := context.Background() + cache, _, recorder := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(102) + embedded := time.Now().UTC().Truncate(time.Millisecond).Add(-time.Hour) + account := memoTestAccount(11, "a") + account.LastUsedAt = &embedded + memoPublish(t, cache, bucket, account) + + first := memoRead(t, cache, bucket) + require.Equal(t, embedded, *first[0].LastUsedAt) + + latest := embedded.Add(42 * time.Second) + require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{account.ID: latest})) + recorder.reset() + second := memoRead(t, cache, bucket) + require.Zero(t, recorder.countPrefix("zrange "), "last_used updates must not invalidate the decoded snapshot") + require.Equal(t, latest, *second[0].LastUsedAt) + require.Equal(t, embedded, *first[0].LastUsedAt, "an earlier result must not observe later reads") + + // 修改返回对象的顶层字段只影响本次结果。 + second[0].Schedulable = false + second[0].Name = "mutated" + second[0].LastUsedAt = nil + second[0].Credentials = map[string]any{"model_mapping": map[string]any{"claude-sonnet": "replaced"}} + require.Equal(t, "replaced", second[0].GetMappedModel("claude-sonnet")) + + // 删除 last_used 侧键（不改精简记录）：若母本曾被套用 last_used，这里会读到 latest。 + require.NoError(t, cache.rdb.Del(ctx, schedulerLastUsedKey("11")).Err()) + recorder.reset() + third := memoRead(t, cache, bucket) + require.Zero(t, recorder.countPrefix("zrange ")) + require.Equal(t, embedded, *third[0].LastUsedAt) + require.True(t, third[0].Schedulable) + require.Equal(t, "a", third[0].Name) + require.Equal(t, "sonnet-a", third[0].GetMappedModel("claude-sonnet")) +} + +func TestSchedulerCacheSnapshotMemoSetAccountIsVisibleImmediately(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(103) + account := memoTestAccount(21, "a") + memoPublish(t, cache, bucket, account, memoTestAccount(22, "b")) + require.True(t, memoRead(t, cache, bucket)[0].Schedulable) + + globalBefore := memoRevision(t, cache, schedulerAccountMetaRevisionKey) + platformBefore := memoRevision(t, cache, memoAnthropicRevisionKey) + account.Schedulable = false + account.Extra = map[string]any{"session_window_utilization": 0.9} + require.NoError(t, cache.SetAccount(ctx, &account)) + require.NotEqual(t, platformBefore, memoRevision(t, cache, memoAnthropicRevisionKey)) + require.Equal(t, globalBefore, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) + + got := memoRead(t, cache, bucket) + require.Equal(t, []int64{21, 22}, memoIDs(got)) + require.False(t, got[0].Schedulable) + require.Equal(t, 0.9, got[0].Extra["session_window_utilization"]) + require.True(t, got[1].Schedulable) +} + +func TestSchedulerCacheSnapshotMemoDeleteAccountMisses(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(104) + memoPublish(t, cache, bucket, memoTestAccount(31, "a"), memoTestAccount(32, "b")) + memoRead(t, cache, bucket) + + revisionBefore := memoRevision(t, cache, schedulerAccountMetaRevisionKey) + require.NoError(t, cache.DeleteAccount(ctx, 32)) + require.NotEqual(t, revisionBefore, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) + memoRequireMiss(t, cache, bucket) + require.Zero(t, memoEntryCount(cache)) +} + +func TestSchedulerCacheSnapshotMemoCrossBucketMetadataRewrite(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucketA := memoTestBucket(105) + bucketB := memoTestBucket(106) + shared := memoTestAccount(41, "shared") + memoPublish(t, cache, bucketA, shared) + memoPublish(t, cache, bucketB, shared, memoTestAccount(42, "b-only")) + require.Equal(t, "shared", memoRead(t, cache, bucketB)[0].Name) + activeB := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucketB)).Val() + + shared.Name = "shared-renamed" + shared.Priority = 9 + memoPublish(t, cache, bucketA, shared) + + require.Equal(t, activeB, cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucketB)).Val()) + got := memoRead(t, cache, bucketB) + require.Equal(t, []int64{41, 42}, memoIDs(got)) + require.Equal(t, "shared-renamed", got[0].Name) + require.Equal(t, 9, got[0].Priority) +} + +func TestSchedulerCacheSnapshotMemoUnencodableLastUsedDeletionMisses(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(107) + memoPublish(t, cache, bucket, memoTestAccount(51, "a"), memoTestAccount(52, "b")) + memoRead(t, cache, bucket) + + revisionBefore := memoRevision(t, cache, schedulerAccountMetaRevisionKey) + invalid := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{52: invalid})) + require.NotEqual(t, revisionBefore, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) + memoRequireMiss(t, cache, bucket) + + // 正常的 last_used 写入不推进任何修订号。 + globalAfterDelete := memoRevision(t, cache, schedulerAccountMetaRevisionKey) + platformAfterDelete := memoRevision(t, cache, memoAnthropicRevisionKey) + require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{51: time.Now()})) + require.Equal(t, globalAfterDelete, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) + require.Equal(t, platformAfterDelete, memoRevision(t, cache, memoAnthropicRevisionKey)) +} + +func TestSchedulerCacheSnapshotMemoMembershipChangeByAccountIDs(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(108) + token, err := cache.CaptureBucketWriteToken(ctx, bucket) + require.NoError(t, err) + ids, err := cache.SetSnapshotAndReturnAccountIDs(ctx, bucket, token, []service.Account{memoTestAccount(61, "a"), memoTestAccount(62, "b")}) + require.NoError(t, err) + require.Equal(t, []int64{61, 62}, ids) + require.Equal(t, []int64{61, 62}, memoIDs(memoRead(t, cache, bucket))) + + globalBefore := memoRevision(t, cache, schedulerAccountMetaRevisionKey) + platformBefore := memoRevision(t, cache, memoAnthropicRevisionKey) + activeBefore := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Val() + token, err = cache.CaptureBucketWriteToken(ctx, bucket) + require.NoError(t, err) + require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, bucket, token, []int64{62})) + + require.Equal(t, globalBefore, memoRevision(t, cache, schedulerAccountMetaRevisionKey), "membership-only publishing does not touch metadata") + require.Equal(t, platformBefore, memoRevision(t, cache, memoAnthropicRevisionKey), "membership-only publishing does not touch metadata") + require.NotEqual(t, activeBefore, cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Val()) + require.Equal(t, []int64{62}, memoIDs(memoRead(t, cache, bucket))) +} + +func TestSchedulerCacheSnapshotMemoRetireReopenAndRepublish(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(109) + memoPublish(t, cache, bucket, memoTestAccount(71, "before")) + memoRead(t, cache, bucket) + + require.NoError(t, cache.RetireBucket(ctx, bucket)) + memoRequireMiss(t, cache, bucket) + require.Zero(t, memoEntryCount(cache)) + + token, err := cache.ReopenBucket(ctx, bucket) + require.NoError(t, err) + memoRequireMiss(t, cache, bucket) + + require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{memoTestAccount(72, "after")})) + got := memoRead(t, cache, bucket) + require.Equal(t, []int64{72}, memoIDs(got)) + require.Equal(t, "after", got[0].Name) +} + +func TestSchedulerCacheSnapshotMemoEmptyAndNilMetadataAreNotCached(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + + t.Run("empty zset", func(t *testing.T) { + bucket := memoTestBucket(110) + _, err := cache.writeAccountIDs(ctx, []service.Account{memoTestAccount(81, "a")}) + require.NoError(t, err) + memoPublish(t, cache, bucket) + memoRequireMiss(t, cache, bucket) + + active := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Val() + require.NotEmpty(t, active) + require.NoError(t, cache.rdb.ZAdd(ctx, schedulerSnapshotKey(bucket, active), redis.Z{Score: 0, Member: "81"}).Err()) + require.Equal(t, []int64{81}, memoIDs(memoRead(t, cache, bucket))) + }) + + t.Run("nil metadata", func(t *testing.T) { + bucket := memoTestBucket(111) + memoPublish(t, cache, bucket, memoTestAccount(91, "a"), memoTestAccount(92, "b")) + metaKey := schedulerAccountMetaKey("92") + payload, err := cache.rdb.Get(ctx, metaKey).Bytes() + require.NoError(t, err) + require.NoError(t, cache.rdb.Del(ctx, metaKey).Err()) + memoRequireMiss(t, cache, bucket) + + require.NoError(t, cache.rdb.Set(ctx, metaKey, payload, 0).Err()) + require.Equal(t, []int64{91, 92}, memoIDs(memoRead(t, cache, bucket))) + }) +} + +func TestSchedulerCacheSnapshotMemoDecodeFailureIsNotCached(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(112) + memoPublish(t, cache, bucket, memoTestAccount(101, "a")) + metaKey := schedulerAccountMetaKey("101") + payload, err := cache.rdb.Get(ctx, metaKey).Bytes() + require.NoError(t, err) + require.NoError(t, cache.rdb.Set(ctx, metaKey, "{not-json", 0).Err()) + + accounts, hit, err := cache.GetSnapshot(ctx, bucket) + require.Error(t, err) + require.False(t, hit) + require.Nil(t, accounts) + require.Zero(t, memoEntryCount(cache)) + + require.NoError(t, cache.rdb.Set(ctx, metaKey, payload, 0).Err()) + require.Equal(t, "a", memoRead(t, cache, bucket)[0].Name) +} + +func TestSchedulerCacheSnapshotMemoRedisErrorIsReturned(t *testing.T) { + cache, mr, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(113) + memoPublish(t, cache, bucket, memoTestAccount(111, "a")) + memoRead(t, cache, bucket) + + mr.SetError("injected failure") + accounts, hit, err := cache.GetSnapshot(context.Background(), bucket) + require.Error(t, err) + require.False(t, hit) + require.Nil(t, accounts) + + mr.SetError("") + require.Equal(t, []int64{111}, memoIDs(memoRead(t, cache, bucket))) +} + +func TestSchedulerCacheSnapshotMemoConcurrentReadersAndWriters(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(114) + accounts := []service.Account{memoTestAccount(121, "a"), memoTestAccount(122, "b"), memoTestAccount(123, "c")} + memoPublish(t, cache, bucket, accounts...) + + var wg sync.WaitGroup + for reader := 0; reader < 8; reader++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 40; i++ { + got, hit, err := cache.GetSnapshot(ctx, bucket) + if err != nil || !hit { + t.Errorf("GetSnapshot: hit=%v err=%v", hit, err) + return + } + for _, account := range got { + if account.GetModelMapping()["claude-sonnet"] == "" { + t.Errorf("account %d lost its model mapping", account.ID) + return + } + _ = account.IsSchedulableForModel("claude-sonnet") + _ = account.GetHeaderOverrides() + account.Schedulable = !account.Schedulable + account.LastUsedAt = nil + } + } + }() + } + for writer := 0; writer < 2; writer++ { + wg.Add(1) + go func(writer int) { + defer wg.Done() + for i := 0; i < 20; i++ { + account := memoTestAccount(int64(121+writer), fmt.Sprintf("w%d-%d", writer, i)) + account.Extra = map[string]any{"session_window_utilization": float64(i) / 100} + if err := cache.SetAccount(ctx, &account); err != nil { + t.Errorf("SetAccount: %v", err) + return + } + if err := cache.UpdateLastUsed(ctx, map[int64]time.Time{account.ID: time.Now()}); err != nil { + t.Errorf("UpdateLastUsed: %v", err) + return + } + } + }(writer) + } + wg.Wait() + + got := memoRead(t, cache, bucket) + require.Equal(t, "w0-19", got[0].Name) + require.Equal(t, "w1-19", got[1].Name) + require.Equal(t, "c", got[2].Name) +} + +func TestSchedulerCacheSnapshotMemoWriteAccountIDsAdvancesRevision(t *testing.T) { + ctx := context.Background() + geminiRevisionKey := schedulerAccountMetaRevisionKeyForPlatform(service.PlatformGemini) + + t.Run("revisions share the last chunk", func(t *testing.T) { + cache, _, recorder := newSchedulerCacheMemoUnit(t, 2) + accounts := []service.Account{ + memoTestAccount(131, "a"), + memoTestAccountOn(service.PlatformGemini, 132, "b"), + memoTestAccount(133, "c"), + } + recorder.reset() + ids, err := cache.writeAccountIDs(ctx, accounts) + require.NoError(t, err) + require.Equal(t, []int64{131, 132, 133}, ids) + + batches := recorder.pipelineBatches() + require.Len(t, batches, 2, "one execution per chunk") + require.Equal(t, []string{ + "set " + schedulerAccountKey("133"), + "set " + schedulerAccountMetaKey("133"), + "set " + memoAnthropicRevisionKey, + "set " + geminiRevisionKey, + }, batches[1]) + require.NotContains(t, strings.Join(batches[0], "\n"), schedulerAccountMetaRevisionKey) + require.Equal(t, 2, recorder.countPrefix("set "+schedulerAccountMetaRevisionKey), "no revision write outside the last chunk") + require.Empty(t, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) + first := memoRevision(t, cache, memoAnthropicRevisionKey) + require.Len(t, first, schedulerAccountMetaRevisionBytes*2) + require.Len(t, memoRevision(t, cache, geminiRevisionKey), schedulerAccountMetaRevisionBytes*2) + + _, err = cache.writeAccountIDs(ctx, accounts[:1]) + require.NoError(t, err) + require.NotEqual(t, first, memoRevision(t, cache, memoAnthropicRevisionKey)) + }) + + t.Run("account count is a multiple of the chunk size", func(t *testing.T) { + cache, _, recorder := newSchedulerCacheMemoUnit(t, 2) + accounts := []service.Account{ + memoTestAccount(141, "a"), memoTestAccount(142, "b"), + memoTestAccount(143, "c"), memoTestAccount(144, "d"), + } + recorder.reset() + _, err := cache.writeAccountIDs(ctx, accounts) + require.NoError(t, err) + + batches := recorder.pipelineBatches() + require.Len(t, batches, 2, "no extra execution for the revision") + require.Equal(t, []string{ + "set " + schedulerAccountKey("143"), + "set " + schedulerAccountMetaKey("143"), + "set " + schedulerAccountKey("144"), + "set " + schedulerAccountMetaKey("144"), + "set " + memoAnthropicRevisionKey, + }, batches[1]) + require.Len(t, recorder.recorded(), 9, "every command travels inside the two pipelines") + }) + + t.Run("failure in a later chunk advances the global revision", func(t *testing.T) { + cache, _, recorder := newSchedulerCacheMemoUnit(t, 1) + accounts := []service.Account{memoTestAccount(151, "a"), memoTestAccount(152, "b"), memoTestAccount(153, "c")} + bucket := memoTestBucket(115) + memoPublish(t, cache, bucket, accounts...) + memoRead(t, cache, bucket) + globalBefore := memoRevision(t, cache, schedulerAccountMetaRevisionKey) + platformBefore := memoRevision(t, cache, memoAnthropicRevisionKey) + + renamed := []service.Account{memoTestAccount(151, "a2"), memoTestAccount(152, "b2"), memoTestAccount(153, "c2")} + recorder.reset() + recorder.setFailPipelineAt(2) + _, err := cache.writeAccountIDs(ctx, renamed) + require.Error(t, err) + recorder.setFailPipelineAt(0) + require.Eventually(t, func() bool { + return memoRevisionValue(cache, schedulerAccountMetaRevisionKey) != globalBefore + }, 2*time.Second, 5*time.Millisecond, "the global revision is advanced in the background") + require.Equal(t, platformBefore, memoRevision(t, cache, memoAnthropicRevisionKey)) + + recorder.reset() + got := memoRead(t, cache, bucket) + require.Equal(t, 1, recorder.countPrefix("zrange "), "the cached snapshot must be reloaded") + require.Equal(t, []string{"a2", "b", "c"}, []string{got[0].Name, got[1].Name, got[2].Name}) + }) + + t.Run("failure in the last chunk advances the global revision", func(t *testing.T) { + cache, _, recorder := newSchedulerCacheMemoUnit(t, 2) + recorder.reset() + recorder.setFailPipelineAt(1) + _, err := cache.writeAccountIDs(ctx, []service.Account{memoTestAccount(161, "a")}) + require.Error(t, err) + recorder.setFailPipelineAt(0) + require.Eventually(t, func() bool { + return memoRevisionValue(cache, schedulerAccountMetaRevisionKey) != "" + }, 2*time.Second, 5*time.Millisecond, "the global revision is advanced in the background") + require.Empty(t, memoRevision(t, cache, memoAnthropicRevisionKey)) + }) + + t.Run("nothing encodable", func(t *testing.T) { + cache, _, recorder := newSchedulerCacheMemoUnit(t, 1) + invalid := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + recorder.reset() + ids, err := cache.writeAccountIDs(ctx, []service.Account{{ID: 171, Platform: service.PlatformAnthropic, ExpiresAt: &invalid}}) + require.NoError(t, err) + require.Empty(t, ids) + require.Empty(t, recorder.recorded()) + require.Empty(t, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) + require.Empty(t, memoRevision(t, cache, memoAnthropicRevisionKey)) + }) + + t.Run("empty platform uses the global revision", func(t *testing.T) { + cache, _, _ := newSchedulerCacheMemoUnit(t, 2) + _, err := cache.writeAccountIDs(ctx, []service.Account{{ID: 181, Name: "no-platform"}}) + require.NoError(t, err) + require.NotEmpty(t, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) + }) +} + +// memoReadExpectHit 读取一次并报告是否重新加载（发出过 ZRANGE）。 +func memoReadExpectHit(t *testing.T, cache *schedulerCache, recorder *schedulerCacheCommandRecorder, bucket service.SchedulerBucket) ([]*service.Account, bool) { + t.Helper() + recorder.reset() + accounts := memoRead(t, cache, bucket) + return accounts, recorder.countPrefix("zrange ") > 0 +} + +func TestSchedulerCacheSnapshotMemoPlatformRevisionScopesInvalidation(t *testing.T) { + ctx := context.Background() + cache, _, recorder := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + geminiSingle := service.SchedulerBucket{GroupID: 201, Platform: service.PlatformGemini, Mode: service.SchedulerModeSingle} + gemini := memoTestAccountOn(service.PlatformGemini, 211, "gemini") + memoPublish(t, cache, geminiSingle, gemini) + _, reloaded := memoReadExpectHit(t, cache, recorder, geminiSingle) + require.True(t, reloaded) + + anthropic := memoTestAccount(212, "anthropic") + require.NoError(t, cache.SetAccount(ctx, &anthropic)) + recorder.reset() + memoRead(t, cache, geminiSingle) + require.Equal(t, []string{ + "mget " + schedulerBucketKey(schedulerReadyPrefix, geminiSingle), + "mget " + schedulerLastUsedKey("211"), + }, recorder.recorded(), "an Anthropic account write must not invalidate a Gemini bucket") + + gemini.Name = "gemini-renamed" + require.NoError(t, cache.SetAccount(ctx, &gemini)) + got, reloaded := memoReadExpectHit(t, cache, recorder, geminiSingle) + require.True(t, reloaded) + require.Equal(t, "gemini-renamed", got[0].Name) +} + +func TestSchedulerCacheSnapshotMemoPlatformChangeInvalidatesOldPlatformBuckets(t *testing.T) { + ctx := context.Background() + cache, _, recorder := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + geminiSingle := service.SchedulerBucket{GroupID: 202, Platform: service.PlatformGemini, Mode: service.SchedulerModeSingle} + moved := memoTestAccountOn(service.PlatformGemini, 221, "moved") + memoPublish(t, cache, geminiSingle, moved) + globalBefore := memoRevision(t, cache, schedulerAccountMetaRevisionKey) + + moved.Name = "same-platform" + require.NoError(t, cache.SetAccount(ctx, &moved)) + require.Equal(t, globalBefore, memoRevision(t, cache, schedulerAccountMetaRevisionKey), "a same-platform rewrite only advances its platform revision") + memoRead(t, cache, geminiSingle) + + moved.Platform = service.PlatformAnthropic + moved.Name = "now-anthropic" + require.NoError(t, cache.SetAccount(ctx, &moved)) + require.NotEqual(t, globalBefore, memoRevision(t, cache, schedulerAccountMetaRevisionKey), "a platform change advances the global revision") + + got, reloaded := memoReadExpectHit(t, cache, recorder, geminiSingle) + require.True(t, reloaded, "the old platform bucket must drop the cached account") + require.Equal(t, service.PlatformAnthropic, got[0].Platform) + require.Equal(t, "now-anthropic", got[0].Name) +} + +func TestSchedulerCacheSnapshotMemoNewAccountOnlyAdvancesPlatformRevision(t *testing.T) { + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + fresh := memoTestAccountOn(service.PlatformGemini, 231, "fresh") + require.NoError(t, cache.SetAccount(context.Background(), &fresh)) + require.Empty(t, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) + require.NotEmpty(t, memoRevision(t, cache, schedulerAccountMetaRevisionKeyForPlatform(service.PlatformGemini))) +} + +func TestSchedulerCacheSnapshotMemoAntigravityRevisionInvalidatesMixedBuckets(t *testing.T) { + ctx := context.Background() + cache, _, recorder := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + antigravity := memoTestAccountOn(service.PlatformAntigravity, 221, "ag") + antigravity.Extra = map[string]any{"mixed_scheduling": true} + gemini := memoTestAccountOn(service.PlatformGemini, 222, "gemini") + anthropic := memoTestAccount(223, "anthropic") + + geminiMixed := service.SchedulerBucket{GroupID: 202, Platform: service.PlatformGemini, Mode: service.SchedulerModeMixed} + anthropicMixed := service.SchedulerBucket{GroupID: 202, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeMixed} + geminiSingle := service.SchedulerBucket{GroupID: 202, Platform: service.PlatformGemini, Mode: service.SchedulerModeSingle} + geminiForced := service.SchedulerBucket{GroupID: 202, Platform: service.PlatformGemini, Mode: service.SchedulerModeForced} + memoPublish(t, cache, geminiMixed, gemini, antigravity) + memoPublish(t, cache, anthropicMixed, anthropic, antigravity) + memoPublish(t, cache, geminiSingle, gemini) + memoPublish(t, cache, geminiForced, gemini) + for _, bucket := range []service.SchedulerBucket{geminiMixed, anthropicMixed, geminiSingle, geminiForced} { + memoRead(t, cache, bucket) + } + + antigravity.Name = "ag-renamed" + require.NoError(t, cache.SetAccount(ctx, &antigravity)) + + got, reloaded := memoReadExpectHit(t, cache, recorder, geminiMixed) + require.True(t, reloaded) + require.Equal(t, "ag-renamed", got[1].Name) + got, reloaded = memoReadExpectHit(t, cache, recorder, anthropicMixed) + require.True(t, reloaded) + require.Equal(t, "ag-renamed", got[1].Name) + _, reloaded = memoReadExpectHit(t, cache, recorder, geminiSingle) + require.False(t, reloaded, "single buckets never contain Antigravity accounts") + _, reloaded = memoReadExpectHit(t, cache, recorder, geminiForced) + require.False(t, reloaded, "forced buckets never contain Antigravity accounts") +} + +func TestSchedulerCacheSnapshotMemoGlobalRevisionInvalidatesEveryPlatform(t *testing.T) { + ctx := context.Background() + buckets := []service.SchedulerBucket{ + {GroupID: 203, Platform: service.PlatformGemini, Mode: service.SchedulerModeSingle}, + {GroupID: 203, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeSingle}, + {GroupID: 203, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}, + } + setup := func(t *testing.T) (*schedulerCache, *schedulerCacheCommandRecorder) { + cache, _, recorder := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + for i, bucket := range buckets { + memoPublish(t, cache, bucket, memoTestAccountOn(bucket.Platform, int64(231+i), bucket.Platform)) + memoRead(t, cache, bucket) + } + // 不属于任何上述 bucket 的账号。 + unrelated := memoTestAccountOn(service.PlatformGrok, 239, "grok") + require.NoError(t, cache.SetAccount(ctx, &unrelated)) + for _, bucket := range buckets { + _, reloaded := memoReadExpectHit(t, cache, recorder, bucket) + require.False(t, reloaded) + } + return cache, recorder + } + + t.Run("DeleteAccount", func(t *testing.T) { + cache, recorder := setup(t) + require.NoError(t, cache.DeleteAccount(ctx, 239)) + for _, bucket := range buckets { + _, reloaded := memoReadExpectHit(t, cache, recorder, bucket) + require.True(t, reloaded, bucket.String()) + } + }) + + t.Run("UpdateLastUsed unencodable deletion", func(t *testing.T) { + cache, recorder := setup(t) + invalid := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{239: invalid})) + for _, bucket := range buckets { + _, reloaded := memoReadExpectHit(t, cache, recorder, bucket) + require.True(t, reloaded, bucket.String()) + } + }) +} + +func TestSchedulerCacheSnapshotMemoFlushAllAndRepublishSameVersion(t *testing.T) { + ctx := context.Background() + cache, mr, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(116) + memoPublish(t, cache, bucket, memoTestAccount(141, "before")) + require.Equal(t, "before", memoRead(t, cache, bucket)[0].Name) + activeBefore := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Val() + + mr.FlushAll() + memoPublish(t, cache, bucket, memoTestAccount(141, "after")) + require.Equal(t, activeBefore, cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Val()) + require.Equal(t, "after", memoRead(t, cache, bucket)[0].Name) +} + +type memoTestClock struct { + mu sync.Mutex + now time.Time +} + +func (c *memoTestClock) Now() time.Time { + c.mu.Lock() + defer c.mu.Unlock() + return c.now +} + +func (c *memoTestClock) Advance(d time.Duration) { + c.mu.Lock() + c.now = c.now.Add(d) + c.mu.Unlock() +} + +func TestSchedulerCacheSnapshotMemoMaxAgeBoundsUnversionedWrites(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + clock := &memoTestClock{now: time.Date(2026, 10, 20, 8, 0, 0, 0, time.UTC)} + cache.now = clock.Now + bucket := memoTestBucket(117) + memoPublish(t, cache, bucket, memoTestAccount(151, "old")) + require.Equal(t, "old", memoRead(t, cache, bucket)[0].Name) + + // 模拟不推进修订号的旧实例：直接改写精简记录。 + _, metaPayload, err := marshalSchedulerCacheAccount(memoTestAccount(151, "new")) + require.NoError(t, err) + globalBefore := memoRevision(t, cache, schedulerAccountMetaRevisionKey) + platformBefore := memoRevision(t, cache, memoAnthropicRevisionKey) + require.NoError(t, cache.rdb.Set(ctx, schedulerAccountMetaKey("151"), metaPayload, 0).Err()) + require.Equal(t, globalBefore, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) + require.Equal(t, platformBefore, memoRevision(t, cache, memoAnthropicRevisionKey)) + + clock.Advance(schedulerSnapshotCacheMaxAge - time.Millisecond) + require.Equal(t, "old", memoRead(t, cache, bucket)[0].Name) + + clock.Advance(time.Millisecond) + require.Equal(t, "new", memoRead(t, cache, bucket)[0].Name) +} + +func TestSchedulerCacheSnapshotMemoEvictsLongIdleEntries(t *testing.T) { + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + clock := &memoTestClock{now: time.Date(2026, 10, 20, 8, 0, 0, 0, time.UTC)} + cache.now = clock.Now + idle := memoTestBucket(118) + recent := memoTestBucket(119) + active := memoTestBucket(120) + memoPublish(t, cache, idle, memoTestAccount(161, "idle")) + memoPublish(t, cache, recent, memoTestAccount(162, "recent")) + memoPublish(t, cache, active, memoTestAccount(163, "active")) + + memoRead(t, cache, idle) + clock.Advance(schedulerSnapshotCacheEvictAge / 2) + memoRead(t, cache, recent) + clock.Advance(schedulerSnapshotCacheEvictAge/2 + time.Second) + memoRead(t, cache, active) + + cache.snapshotMu.RLock() + defer cache.snapshotMu.RUnlock() + require.NotContains(t, cache.snapshotEntries, idle.String()) + require.Contains(t, cache.snapshotEntries, recent.String()) + require.Contains(t, cache.snapshotEntries, active.String()) +} + +func TestSchedulerCacheSnapshotMemoKeepsNewerEntryOnLateStore(t *testing.T) { + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + base := time.Date(2026, 10, 20, 8, 0, 0, 0, time.UTC) + newer := &schedulerSnapshotEntry{active: "2", revisions: []string{"r2"}, loadedAt: base.Add(time.Second)} + older := &schedulerSnapshotEntry{active: "1", revisions: []string{"r1"}, loadedAt: base} + cache.storeSnapshotEntry("b", newer) + cache.storeSnapshotEntry("b", older) + require.Same(t, newer, cache.lookupSnapshotEntry("b")) +} + +// 母本在发布前预解析了模型映射：副本命中时不计算签名。此处故意原地改写共享的原始映射， +// 只用来观察副本走的是免签名路径（生产代码不会原地修改该 map）。 +func TestSchedulerCacheSnapshotMemoPreparesSharedModelMapping(t *testing.T) { + cache, _, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + bucket := memoTestBucket(121) + memoPublish(t, cache, bucket, memoTestAccount(171, "a")) + memoRead(t, cache, bucket) + + hit := memoRead(t, cache, bucket)[0] + require.Equal(t, "sonnet-a", hit.GetMappedModel("claude-sonnet")) + + replaced := memoRead(t, cache, bucket)[0] + replaced.Credentials = map[string]any{"model_mapping": map[string]any{"claude-sonnet": "sonnet-replaced"}} + require.Equal(t, "sonnet-replaced", replaced.GetMappedModel("claude-sonnet")) + + probe := memoRead(t, cache, bucket)[0] + rawMapping, ok := probe.Credentials["model_mapping"].(map[string]any) + require.True(t, ok) + rawMapping["claude-sonnet"] = "sonnet-mutated-in-place" + require.Equal(t, "sonnet-a", probe.GetMappedModel("claude-sonnet")) +} + +func TestSchedulerCacheSnapshotMemoRevisionValuesAreUnique(t *testing.T) { + seen := make(map[string]struct{}, 64) + for i := 0; i < 64; i++ { + revision := newSchedulerAccountMetaRevision() + require.Len(t, revision, schedulerAccountMetaRevisionBytes*2) + _, err := strconv.ParseUint(revision[:8], 16, 32) + require.NoError(t, err) + _, dup := seen[revision] + require.False(t, dup) + seen[revision] = struct{}{} + } +} + +func TestSchedulerCacheSnapshotMemoPlatformChangeInEarlierChunk(t *testing.T) { + ctx := context.Background() + cache, _, _ := newSchedulerCacheMemoUnit(t, 1) + require.NoError(t, cache.SetAccount(ctx, ptrAccount(memoTestAccountOn(service.PlatformGemini, 241, "moved")))) + require.NoError(t, cache.SetAccount(ctx, ptrAccount(memoTestAccount(242, "stays")))) + globalBefore := memoRevision(t, cache, schedulerAccountMetaRevisionKey) + + _, err := cache.writeAccountIDs(ctx, []service.Account{ + memoTestAccount(241, "moved"), + memoTestAccount(242, "stays"), + }) + require.NoError(t, err) + require.NotEqual(t, globalBefore, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) +} + +func TestSchedulerCacheSnapshotMemoUnparsableOldMetadataAdvancesGlobalRevision(t *testing.T) { + ctx := context.Background() + cache, mr, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + require.NoError(t, mr.Set(schedulerAccountMetaKey("251"), "not-json")) + + require.NoError(t, cache.SetAccount(ctx, ptrAccount(memoTestAccount(251, "replaced")))) + require.NotEmpty(t, memoRevision(t, cache, schedulerAccountMetaRevisionKey)) +} + +func TestSchedulerCacheSnapshotMemoNilBeforeRealPipelineError(t *testing.T) { + ctx := context.Background() + cache, mr, _ := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + mr.HSet(schedulerAccountMetaKey("262"), "field", "value") + + _, err := cache.writeAccountIDs(ctx, []service.Account{ + memoTestAccount(261, "fresh"), + memoTestAccount(262, "wrong-type"), + }) + require.Error(t, err, "a WRONGTYPE reply after a nil SET ... GET reply must fail the write") + require.Eventually(t, func() bool { + return memoRevisionValue(cache, schedulerAccountMetaRevisionKey) != "" + }, 2*time.Second, 5*time.Millisecond, "the global revision is advanced in the background") +} + +func TestSchedulerCacheSnapshotMemoCoalescesBackgroundGlobalRevisionWrites(t *testing.T) { + cache, _, recorder := newSchedulerCacheMemoUnit(t, defaultSchedulerSnapshotWriteChunkSize) + recorder.reset() + for range 50 { + cache.scheduleGlobalAccountMetaRevision(context.Background()) + } + require.Eventually(t, func() bool { + return !cache.globalRevisionWriting.Load() && memoRevisionValue(cache, schedulerAccountMetaRevisionKey) != "" + }, 2*time.Second, 5*time.Millisecond) + writes := 0 + for _, command := range recorder.recorded() { + if command == "set "+schedulerAccountMetaRevisionKey { + writes++ + } + } + require.GreaterOrEqual(t, writes, 1) + require.LessOrEqual(t, writes, 3, "concurrent schedules collapse into at most a couple of background writes") +} + +func ptrAccount(account service.Account) *service.Account { + return &account +} diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index 4e01e4195922..1126d2ddc35f 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -33,148 +33,10 @@ func newSchedulerCacheUnitWithRedis(t *testing.T) (*schedulerCache, *miniredis.M return cache, mr } -func TestSchedulerCacheIgnoresLegacyMetadataGeneration(t *testing.T) { - ctx := context.Background() - cache, _ := newSchedulerCacheUnitWithRedis(t) - bucket := service.SchedulerBucket{GroupID: 91, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} - account := service.Account{ - ID: 9101, - Platform: service.PlatformOpenAI, - Type: service.AccountTypeAPIKey, - Status: service.StatusActive, - Credentials: map[string]any{ - "api_key": "test-only-key", - "base_url": "https://relay.example.test", - }, - } - legacy := buildSchedulerMetadataAccount(account) - delete(legacy.Credentials, "base_url") - delete(legacy.Credentials, "openai_capabilities") - legacyJSON, err := json.Marshal(legacy) - require.NoError(t, err) - - const legacyVersion = "1" - require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey(schedulerLegacyReadyPrefix, bucket), "1", 0).Err()) - require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey("sched:active:", bucket), legacyVersion, 0).Err()) - require.NoError(t, cache.rdb.ZAdd(ctx, fmt.Sprintf("sched:%d:%s:%s:v%s", bucket.GroupID, bucket.Platform, bucket.Mode, legacyVersion), redis.Z{ - Score: 0, - Member: strconv.FormatInt(account.ID, 10), - }).Err()) - require.NoError(t, cache.rdb.Set(ctx, schedulerLegacyAccountMetaKey(strconv.FormatInt(account.ID, 10)), legacyJSON, 0).Err()) - - legacySnapshot, hit, err := cache.GetSnapshot(ctx, bucket) - require.NoError(t, err) - require.False(t, hit) - require.Nil(t, legacySnapshot) - - token, err := cache.CaptureBucketWriteToken(ctx, bucket) - require.NoError(t, err) - require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account})) - - currentSnapshot, hit, err := cache.GetSnapshot(ctx, bucket) - require.NoError(t, err) - require.True(t, hit) - require.Len(t, currentSnapshot, 1) - require.Equal(t, "https://relay.example.test", currentSnapshot[0].Credentials["base_url"]) -} - -func TestSchedulerCacheMetadataGenerationCleansLegacyCredentialCopies(t *testing.T) { - ctx := context.Background() - cache := newSchedulerCacheUnit(t) - const accountID int64 = 9102 - id := strconv.FormatInt(accountID, 10) - legacyMetaKey := "sched:meta:" + id - - require.NoError(t, cache.rdb.Set(ctx, legacyMetaKey, `{"credentials":{"api_key":"retired-key"}}`, 0).Err()) - require.NoError(t, cache.SetAccount(ctx, &service.Account{ - ID: accountID, - Platform: service.PlatformOpenAI, - Type: service.AccountTypeAPIKey, - Status: service.StatusActive, - Schedulable: true, - Credentials: map[string]any{"api_key": "rotated-key"}, - })) - - _, err := cache.rdb.Get(ctx, legacyMetaKey).Result() - require.ErrorIs(t, err, redis.Nil, "credential rotation must remove the legacy metadata copy") - require.True(t, cache.rdb.Exists(ctx, schedulerAccountMetaKey(id)).Val() > 0) - - require.NoError(t, cache.rdb.Set(ctx, legacyMetaKey, `{"credentials":{"api_key":"retired-key"}}`, 0).Err()) - require.NoError(t, cache.DeleteAccount(ctx, accountID)) - require.Zero(t, cache.rdb.Exists(ctx, - schedulerAccountKey(id), - schedulerAccountMetaKey(id), - legacyMetaKey, - schedulerLastUsedKey(id), - ).Val(), "account deletion must remove both metadata generations") -} - -func TestSchedulerCacheV2ActivationKeepsLegacyReaderGenerationConsistent(t *testing.T) { - ctx := context.Background() - cache := newSchedulerCacheUnit(t) - bucket := service.SchedulerBucket{GroupID: 92, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} - legacyAccount := service.Account{ - ID: 9201, - Platform: service.PlatformOpenAI, - Type: service.AccountTypeAPIKey, - Status: service.StatusActive, - Schedulable: true, - Credentials: map[string]any{"api_key": "legacy-key"}, - } - currentAccount := service.Account{ - ID: 9202, - Platform: service.PlatformOpenAI, - Type: service.AccountTypeAPIKey, - Status: service.StatusActive, - Schedulable: true, - Credentials: map[string]any{ - "api_key": "current-key", - "base_url": "https://api.laxarouter.ai", - }, - } - legacyPayload, err := json.Marshal(buildSchedulerMetadataAccount(legacyAccount)) - require.NoError(t, err) - staleCurrent := buildSchedulerMetadataAccount(currentAccount) - delete(staleCurrent.Credentials, "base_url") - staleCurrentPayload, err := json.Marshal(staleCurrent) - require.NoError(t, err) - - const legacyVersion = "41" - legacyReadyKey := schedulerBucketKey("sched:ready:", bucket) - legacyActiveKey := schedulerBucketKey("sched:active:", bucket) - legacySnapshotKey := fmt.Sprintf("sched:%d:%s:%s:v%s", bucket.GroupID, bucket.Platform, bucket.Mode, legacyVersion) - require.NoError(t, cache.rdb.Set(ctx, legacyReadyKey, "1", 0).Err()) - require.NoError(t, cache.rdb.Set(ctx, legacyActiveKey, legacyVersion, 0).Err()) - require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey("sched:ver:", bucket), legacyVersion, 0).Err()) - require.NoError(t, cache.rdb.ZAdd(ctx, legacySnapshotKey, redis.Z{ - Score: 0, - Member: strconv.FormatInt(legacyAccount.ID, 10), - }).Err()) - require.NoError(t, cache.rdb.Set(ctx, "sched:meta:"+strconv.FormatInt(legacyAccount.ID, 10), legacyPayload, 0).Err()) - require.NoError(t, cache.rdb.Set(ctx, "sched:meta:"+strconv.FormatInt(currentAccount.ID, 10), staleCurrentPayload, 0).Err()) - - token, err := cache.CaptureBucketWriteToken(ctx, bucket) - require.NoError(t, err) - require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{currentAccount})) - - _, err = cache.rdb.Get(ctx, legacyReadyKey).Result() - require.ErrorIs(t, err, redis.Nil, "v2 activation must force a rolled-back reader through DB fallback") - require.Equal(t, legacyVersion, cache.rdb.Get(ctx, legacyActiveKey).Val(), - "v2 activation must not point a legacy reader at v2 snapshot members") - require.Equal(t, []string{strconv.FormatInt(legacyAccount.ID, 10)}, cache.rdb.ZRange(ctx, legacySnapshotKey, 0, -1).Val()) - - currentSnapshot, hit, err := cache.GetSnapshot(ctx, bucket) - require.NoError(t, err) - require.True(t, hit) - require.Len(t, currentSnapshot, 1) - require.Equal(t, currentAccount.ID, currentSnapshot[0].ID) -} - func TestSchedulerCacheWriteAccountIDsSkipsUnencodableTimes(t *testing.T) { ctx := context.Background() cache := newSchedulerCacheUnit(t) invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) - require.NoError(t, cache.rdb.Set(ctx, schedulerLegacyAccountMetaKey("112"), `{"credentials":{"api_key":"retired-key"}}`, 0).Err()) accountIDs, err := cache.writeAccountIDs(ctx, []service.Account{ {ID: 111, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}, @@ -190,7 +52,6 @@ func TestSchedulerCacheWriteAccountIDsSkipsUnencodableTimes(t *testing.T) { invalid, err := cache.GetAccount(ctx, 112) require.NoError(t, err) require.Nil(t, invalid) - require.Zero(t, cache.rdb.Exists(ctx, schedulerLegacyAccountMetaKey("112")).Val()) } func TestSchedulerCacheSetAccountClearsUnencodablePayload(t *testing.T) { @@ -199,7 +60,6 @@ func TestSchedulerCacheSetAccountClearsUnencodablePayload(t *testing.T) { account := service.Account{ID: 113, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey} require.NoError(t, cache.SetAccount(ctx, &account)) - require.NoError(t, cache.rdb.Set(ctx, schedulerLegacyAccountMetaKey("113"), `{"credentials":{"api_key":"retired-key"}}`, 0).Err()) invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) account.ExpiresAt = &invalidTime @@ -208,7 +68,6 @@ func TestSchedulerCacheSetAccountClearsUnencodablePayload(t *testing.T) { cached, err := cache.GetAccount(ctx, account.ID) require.NoError(t, err) require.Nil(t, cached) - require.Zero(t, cache.rdb.Exists(ctx, schedulerLegacyAccountMetaKey("113")).Val()) } func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) { @@ -216,7 +75,6 @@ func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) { cache := newSchedulerCacheUnit(t) account := service.Account{ID: 114, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey} require.NoError(t, cache.SetAccount(ctx, &account)) - require.NoError(t, cache.rdb.Set(ctx, schedulerLegacyAccountMetaKey("114"), `{"credentials":{"api_key":"retired-key"}}`, 0).Err()) invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{account.ID: invalidTime})) @@ -224,7 +82,6 @@ func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) { cached, err := cache.GetAccount(ctx, account.ID) require.NoError(t, err) require.Nil(t, cached) - require.Zero(t, cache.rdb.Exists(ctx, schedulerLegacyAccountMetaKey("114")).Val()) } func TestSchedulerCacheSnapshotAccountIDReusePreservesPayloadAndMembers(t *testing.T) { diff --git a/backend/internal/repository/simple_mode_default_groups.go b/backend/internal/repository/simple_mode_default_groups.go index 2b9daa4d6411..25f3d6df5f95 100644 --- a/backend/internal/repository/simple_mode_default_groups.go +++ b/backend/internal/repository/simple_mode_default_groups.go @@ -77,7 +77,7 @@ func createGroupIfNotExists(ctx context.Context, client *dbent.Client, name, pla SetSubscriptionType(service.SubscriptionTypeStandard). SetRateMultiplier(1.0). SetIsExclusive(false). - SetAllowImageGeneration(platform == service.PlatformGrok). + SetAllowImageGeneration(service.DefaultAllowImageGenerationForPlatform(platform)). Save(ctx) if err != nil { if dbent.IsConstraintError(err) { diff --git a/backend/internal/repository/simple_mode_default_groups_integration_test.go b/backend/internal/repository/simple_mode_default_groups_integration_test.go index d6a96bbfa248..e21d1b0b4d86 100644 --- a/backend/internal/repository/simple_mode_default_groups_integration_test.go +++ b/backend/internal/repository/simple_mode_default_groups_integration_test.go @@ -34,11 +34,18 @@ func TestEnsureSimpleModeDefaultGroups_CreatesMissingDefaults(t *testing.T) { assertGroupExists(service.PlatformAntigravity + "-default-1") assertGroupExists(service.PlatformAntigravity + "-default-2") - grokDefault, err := client.Group.Query(). - Where(group.NameEQ(service.PlatformGrok+"-default"), group.DeletedAtIsNil()). - Only(seedCtx) - require.NoError(t, err) - require.True(t, grokDefault.AllowImageGeneration) + for _, name := range []string{ + service.PlatformGrok + "-default", + service.PlatformGemini + "-default", + service.PlatformAntigravity + "-default-1", + service.PlatformAntigravity + "-default-2", + } { + imageDefault, err := client.Group.Query(). + Where(group.NameEQ(name), group.DeletedAtIsNil()). + Only(seedCtx) + require.NoError(t, err) + require.True(t, imageDefault.AllowImageGeneration, "expected group %s to allow image generation", name) + } } func TestEnsureSimpleModeDefaultGroups_BackfillsOnlyAutoCreatedGrokDefault(t *testing.T) { diff --git a/backend/internal/repository/usage_billing_repo.go b/backend/internal/repository/usage_billing_repo.go index bfb1bb0de5f4..6b509215b301 100644 --- a/backend/internal/repository/usage_billing_repo.go +++ b/backend/internal/repository/usage_billing_repo.go @@ -202,9 +202,10 @@ func (r *usageBillingRepository) applyUsageBillingEffects(ctx context.Context, t } } - if cmd.AccountQuotaCost > 0 && (strings.EqualFold(cmd.AccountType, service.AccountTypeAPIKey) || strings.EqualFold(cmd.AccountType, service.AccountTypeBedrock)) { - quotaState, err := incrementUsageBillingAccountQuota(ctx, tx, cmd.AccountID, cmd.AccountQuotaCost) - if err != nil { + // 账号已不存在时跳过其配额计数，其余结算项不受影响。 + if cmd.AccountQuotaCost > 0 && service.AccountTypeSupportsQuota(cmd.AccountType) { + quotaState, err := incrementUsageBillingAccountQuota(ctx, tx, cmd.AccountID, cmd.AccountQuotaCost, cmd.AccountQuotaHasDailyLimit, cmd.AccountQuotaHasWeeklyLimit) + if err != nil && !errors.Is(err, service.ErrAccountNotFound) { return err } result.QuotaState = quotaState @@ -465,57 +466,157 @@ func incrementUsageBillingAPIKeyRateLimit(ctx context.Context, tx *sql.Tx, apiKe return nil } -func incrementUsageBillingAccountQuota(ctx context.Context, tx *sql.Tx, accountID int64, amount float64) (*service.AccountQuotaState, error) { - rows, err := tx.QueryContext(ctx, - `UPDATE accounts SET extra = ( - COALESCE(extra, '{}'::jsonb) - || jsonb_build_object('quota_used', COALESCE((extra->>'quota_used')::numeric, 0) + $1) +// accountQuotaTotalIncrementExpr 累加总额度已用量。 +const accountQuotaTotalIncrementExpr = `COALESCE(extra, '{}'::jsonb) + || jsonb_build_object('quota_used', COALESCE((extra->>'quota_used')::numeric, 0) + $1)` + +// accountQuotaDailyIncrementExpr 在库中设置了日上限时累加日用量，周期过期时从本次金额重新计数。 +const accountQuotaDailyIncrementExpr = ` || CASE WHEN COALESCE((extra->>'quota_daily_limit')::numeric, 0) > 0 THEN jsonb_build_object( 'quota_daily_used', - CASE WHEN `+dailyExpiredExpr+` + CASE WHEN ` + dailyExpiredExpr + ` THEN $1 ELSE COALESCE((extra->>'quota_daily_used')::numeric, 0) + $1 END, 'quota_daily_start', - CASE WHEN `+dailyExpiredExpr+` - THEN `+nowUTC+` - ELSE COALESCE(extra->>'quota_daily_start', `+nowUTC+`) END + CASE WHEN ` + dailyExpiredExpr + ` + THEN ` + nowUTC + ` + ELSE COALESCE(extra->>'quota_daily_start', ` + nowUTC + `) END ) - || CASE WHEN `+dailyExpiredExpr+` AND `+nextDailyResetAtExpr+` IS NOT NULL - THEN jsonb_build_object('quota_daily_reset_at', `+nextDailyResetAtExpr+`) + || CASE WHEN ` + dailyExpiredExpr + ` AND ` + nextDailyResetAtExpr + ` IS NOT NULL + THEN jsonb_build_object('quota_daily_reset_at', ` + nextDailyResetAtExpr + `) ELSE '{}'::jsonb END - ELSE '{}'::jsonb END + ELSE '{}'::jsonb END` + +// accountQuotaWeeklyIncrementExpr 在库中设置了周上限时累加周用量，周期过期时从本次金额重新计数。 +const accountQuotaWeeklyIncrementExpr = ` || CASE WHEN COALESCE((extra->>'quota_weekly_limit')::numeric, 0) > 0 THEN jsonb_build_object( 'quota_weekly_used', - CASE WHEN `+weeklyExpiredExpr+` + CASE WHEN ` + weeklyExpiredExpr + ` THEN $1 ELSE COALESCE((extra->>'quota_weekly_used')::numeric, 0) + $1 END, 'quota_weekly_start', - CASE WHEN `+weeklyExpiredExpr+` - THEN `+nowUTC+` - ELSE COALESCE(extra->>'quota_weekly_start', `+nowUTC+`) END + CASE WHEN ` + weeklyExpiredExpr + ` + THEN ` + nowUTC + ` + ELSE COALESCE(extra->>'quota_weekly_start', ` + nowUTC + `) END ) - || CASE WHEN `+weeklyExpiredExpr+` AND `+nextWeeklyResetAtExpr+` IS NOT NULL - THEN jsonb_build_object('quota_weekly_reset_at', `+nextWeeklyResetAtExpr+`) + || CASE WHEN ` + weeklyExpiredExpr + ` AND ` + nextWeeklyResetAtExpr + ` IS NOT NULL + THEN jsonb_build_object('quota_weekly_reset_at', ` + nextWeeklyResetAtExpr + `) ELSE '{}'::jsonb END - ELSE '{}'::jsonb END - ), updated_at = NOW() - WHERE id = $2 AND deleted_at IS NULL + ELSE '{}'::jsonb END` + +// accountQuotaExtraIsObjectCond 要求 extra 为 JSON 对象：省略某个维度的语句只在对象上与完整语句等价。 +// +// accountQuotaDailyLimitUnsetCond / accountQuotaWeeklyLimitUnsetCond 在库中未设置日/周上限时为真： +// 键不存在、值为 JSON null 或不大于 0 的数字。其他取值一律为假，交给包含该维度的语句处理。 +// 三个条件都只使用 ->、jsonb_typeof 与 jsonb 比较，对任何行求值都不会报错。 +const ( + accountQuotaExtraIsObjectCond = `jsonb_typeof(extra) = 'object'` + accountQuotaDailyLimitUnsetCond = `(extra->'quota_daily_limit' IS NULL + OR jsonb_typeof(extra->'quota_daily_limit') = 'null' + OR (jsonb_typeof(extra->'quota_daily_limit') = 'number' AND extra->'quota_daily_limit' <= '0'::jsonb))` + accountQuotaWeeklyLimitUnsetCond = `(extra->'quota_weekly_limit' IS NULL + OR jsonb_typeof(extra->'quota_weekly_limit') = 'null' + OR (jsonb_typeof(extra->'quota_weekly_limit') = 'number' AND extra->'quota_weekly_limit' <= '0'::jsonb))` +) + +const accountQuotaIncrementReturning = ` RETURNING COALESCE((extra->>'quota_used')::numeric, 0), COALESCE((extra->>'quota_limit')::numeric, 0), COALESCE((extra->>'quota_daily_used')::numeric, 0), COALESCE((extra->>'quota_daily_limit')::numeric, 0), COALESCE((extra->>'quota_weekly_used')::numeric, 0), - COALESCE((extra->>'quota_weekly_limit')::numeric, 0)`, - amount, accountID) + COALESCE((extra->>'quota_weekly_limit')::numeric, 0)` + +// accountQuotaIncrementSQLs 按 accountQuotaIncrementSQLIndex 存放账号配额累加语句。省略日或周维度的语句 +// 在 WHERE 中附加 accountQuotaExtraIsObjectCond 与被省略维度的未设置条件，条件不成立时命中 0 行； +// 同时包含日、周维度的语句不附加这些条件，对任何账号配置都给出完整的累加结果。 +var accountQuotaIncrementSQLs = buildAccountQuotaIncrementSQLs() + +func accountQuotaIncrementSQLIndex(daily, weekly bool) int { + index := 0 + if daily { + index |= 1 + } + if weekly { + index |= 2 + } + return index +} + +func buildAccountQuotaIncrementSQLs() [4]string { + var sqls [4]string + for _, daily := range []bool{false, true} { + for _, weekly := range []bool{false, true} { + set := accountQuotaTotalIncrementExpr + where := "id = $2 AND deleted_at IS NULL" + if !daily || !weekly { + where += "\n\t\t\tAND " + accountQuotaExtraIsObjectCond + } + if daily { + set += accountQuotaDailyIncrementExpr + } else { + where += "\n\t\t\tAND " + accountQuotaDailyLimitUnsetCond + } + if weekly { + set += accountQuotaWeeklyIncrementExpr + } else { + where += "\n\t\t\tAND " + accountQuotaWeeklyLimitUnsetCond + } + sqls[accountQuotaIncrementSQLIndex(daily, weekly)] = "UPDATE accounts SET extra = (\n\t\t\t" + set + + "\n\t\t), updated_at = NOW()\n\t\tWHERE " + where + accountQuotaIncrementReturning + } + } + return sqls +} + +// incrementUsageBillingAccountQuota 累加账号配额并返回累加后的状态。hasDaily/hasWeekly 是调用方所见的 +// 账号日/周上限设置，只决定先执行哪条语句；该语句的附加条件不成立而命中 0 行时，再执行同时包含日、周维度的语句。 +func incrementUsageBillingAccountQuota(ctx context.Context, tx *sql.Tx, accountID int64, amount float64, hasDaily, hasWeekly bool) (*service.AccountQuotaState, error) { + index := accountQuotaIncrementSQLIndex(hasDaily, hasWeekly) + fullIndex := accountQuotaIncrementSQLIndex(true, true) + state, err := execAccountQuotaIncrement(ctx, tx, accountQuotaIncrementSQLs[index], accountID, amount) + if err != nil { + return nil, err + } + if state == nil && index != fullIndex { + state, err = execAccountQuotaIncrement(ctx, tx, accountQuotaIncrementSQLs[fullIndex], accountID, amount) + if err != nil { + return nil, err + } + } + if state == nil { + return nil, service.ErrAccountNotFound + } + // 任意维度额度在本次递增中从"未超"跨越到"已超"时，必须刷新调度快照， + // 否则 Redis 中缓存的 Account 仍显示旧的 used 值，后续请求会继续选中本账号， + // 最终观察到 daily_used / weekly_used 大幅超过配置的 limit。 + // 对于日/周额度，即使本次触发了周期重置（pre=0、post=amount）， + // 判定式 (post-amount) < limit 同样成立，逻辑与总额度保持一致。 + crossedTotal := state.TotalLimit > 0 && state.TotalUsed >= state.TotalLimit && (state.TotalUsed-amount) < state.TotalLimit + crossedDaily := state.DailyLimit > 0 && state.DailyUsed >= state.DailyLimit && (state.DailyUsed-amount) < state.DailyLimit + crossedWeekly := state.WeeklyLimit > 0 && state.WeeklyUsed >= state.WeeklyLimit && (state.WeeklyUsed-amount) < state.WeeklyLimit + if crossedTotal || crossedDaily || crossedWeekly { + if err := enqueueSchedulerOutbox(ctx, tx, service.SchedulerOutboxEventAccountChanged, &accountID, nil, nil); err != nil { + logger.LegacyPrintf("repository.usage_billing", "[SchedulerOutbox] enqueue quota exceeded failed: account=%d err=%v", accountID, err) + return nil, err + } + } + return state, nil +} + +// execAccountQuotaIncrement 执行一条账号配额累加语句，未命中任何行时返回 nil 状态。 +func execAccountQuotaIncrement(ctx context.Context, tx *sql.Tx, query string, accountID int64, amount float64) (*service.AccountQuotaState, error) { + rows, err := tx.QueryContext(ctx, query, amount, accountID) if err != nil { return nil, err } - var state service.AccountQuotaState + var state *service.AccountQuotaState if rows.Next() { + state = &service.AccountQuotaState{} if err := rows.Scan( &state.TotalUsed, &state.TotalLimit, &state.DailyUsed, &state.DailyLimit, @@ -524,13 +625,6 @@ func incrementUsageBillingAccountQuota(ctx context.Context, tx *sql.Tx, accountI _ = rows.Close() return nil, err } - } else { - if err := rows.Err(); err != nil { - _ = rows.Close() - return nil, err - } - _ = rows.Close() - return nil, service.ErrAccountNotFound } if err := rows.Err(); err != nil { _ = rows.Close() @@ -542,19 +636,5 @@ func incrementUsageBillingAccountQuota(ctx context.Context, tx *sql.Tx, accountI if err := rows.Close(); err != nil { return nil, err } - // 任意维度额度在本次递增中从"未超"跨越到"已超"时，必须刷新调度快照， - // 否则 Redis 中缓存的 Account 仍显示旧的 used 值，后续请求会继续选中本账号， - // 最终观察到 daily_used / weekly_used 大幅超过配置的 limit。 - // 对于日/周额度，即使本次触发了周期重置（pre=0、post=amount）， - // 判定式 (post-amount) < limit 同样成立，逻辑与总额度保持一致。 - crossedTotal := state.TotalLimit > 0 && state.TotalUsed >= state.TotalLimit && (state.TotalUsed-amount) < state.TotalLimit - crossedDaily := state.DailyLimit > 0 && state.DailyUsed >= state.DailyLimit && (state.DailyUsed-amount) < state.DailyLimit - crossedWeekly := state.WeeklyLimit > 0 && state.WeeklyUsed >= state.WeeklyLimit && (state.WeeklyUsed-amount) < state.WeeklyLimit - if crossedTotal || crossedDaily || crossedWeekly { - if err := enqueueSchedulerOutbox(ctx, tx, service.SchedulerOutboxEventAccountChanged, &accountID, nil, nil); err != nil { - logger.LegacyPrintf("repository.usage_billing", "[SchedulerOutbox] enqueue quota exceeded failed: account=%d err=%v", accountID, err) - return nil, err - } - } - return &state, nil + return state, nil } diff --git a/backend/internal/repository/usage_billing_repo_integration_test.go b/backend/internal/repository/usage_billing_repo_integration_test.go index 9f2f0169dcec..9ca7ea03ceec 100644 --- a/backend/internal/repository/usage_billing_repo_integration_test.go +++ b/backend/internal/repository/usage_billing_repo_integration_test.go @@ -5,11 +5,13 @@ package repository import ( "context" "fmt" + "maps" "strings" "testing" "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgconn" "github.com/stretchr/testify/require" "github.com/Wei-Shaw/sub2api/internal/service" @@ -217,6 +219,54 @@ func TestUsageBillingRepositoryApply_DeletedAPIKeyStillBillsBalance(t *testing.T require.Equal(t, 1, dedupCount) } +func TestUsageBillingRepositoryApply_DeletedAccountStillBillsBalance(t *testing.T) { + ctx := context.Background() + client := testEntClient(t) + repo := NewUsageBillingRepository(client, integrationDB) + + user := mustCreateUser(t, client, &service.User{ + Email: fmt.Sprintf("usage-billing-deleted-account-user-%d@example.com", time.Now().UnixNano()), + PasswordHash: "hash", + Balance: 100, + }) + apiKey := mustCreateApiKey(t, client, &service.APIKey{ + UserID: user.ID, + Key: "sk-usage-billing-deleted-account-" + uuid.NewString(), + Name: "billing-deleted-account", + }) + account := mustCreateAccount(t, client, &service.Account{ + Name: "usage-billing-deleted-account-" + uuid.NewString(), + Platform: service.PlatformGemini, + Type: service.AccountTypeServiceAccount, + }) + + _, err := integrationDB.ExecContext(ctx, "UPDATE accounts SET deleted_at = NOW() WHERE id = $1", account.ID) + require.NoError(t, err) + + requestID := uuid.NewString() + result, err := repo.Apply(ctx, &service.UsageBillingCommand{ + RequestID: requestID, + APIKeyID: apiKey.ID, + UserID: user.ID, + AccountID: account.ID, + AccountType: service.AccountTypeServiceAccount, + BalanceCost: 1.25, + AccountQuotaCost: 1.25, + }) + require.NoError(t, err) + require.NotNil(t, result) + require.True(t, result.Applied) + require.Nil(t, result.QuotaState) + + var balance float64 + require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT balance FROM users WHERE id = $1", user.ID).Scan(&balance)) + require.InDelta(t, 98.75, balance, 0.000001) + + var dedupCount int + require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1 AND api_key_id = $2", requestID, apiKey.ID).Scan(&dedupCount)) + require.Equal(t, 1, dedupCount) +} + func TestUsageBillingRepositoryApply_UpdatesAccountQuota(t *testing.T) { ctx := context.Background() client := testEntClient(t) @@ -254,6 +304,60 @@ func TestUsageBillingRepositoryApply_UpdatesAccountQuota(t *testing.T) { require.InDelta(t, 3.5, quotaUsed, 0.000001) } +// service_account 只设总上限：累加 quota_used；日/周用量只在设置了对应 limit 时累加，未跨越上限不入队调度 outbox。 +func TestUsageBillingRepositoryApply_UpdatesServiceAccountQuota(t *testing.T) { + ctx := context.Background() + client := testEntClient(t) + repo := NewUsageBillingRepository(client, integrationDB) + + user := mustCreateUser(t, client, &service.User{ + Email: fmt.Sprintf("usage-billing-sa-user-%d@example.com", time.Now().UnixNano()), + PasswordHash: "hash", + }) + apiKey := mustCreateApiKey(t, client, &service.APIKey{ + UserID: user.ID, + Key: "sk-usage-billing-sa-" + uuid.NewString(), + Name: "billing-sa", + }) + account := mustCreateAccount(t, client, &service.Account{ + Name: "usage-billing-sa-quota-" + uuid.NewString(), + Platform: service.PlatformGemini, + Type: service.AccountTypeServiceAccount, + Extra: map[string]any{ + "quota_limit": 100.0, + }, + }) + + for _, cost := range []float64{2.25, 1.5} { + _, err := repo.Apply(ctx, &service.UsageBillingCommand{ + RequestID: uuid.NewString(), + APIKeyID: apiKey.ID, + UserID: user.ID, + AccountID: account.ID, + AccountType: service.AccountTypeServiceAccount, + AccountQuotaCost: cost, + }) + require.NoError(t, err) + } + + var quotaUsed float64 + var hasDailyUsed, hasWeeklyUsed bool + require.NoError(t, integrationDB.QueryRowContext(ctx, + "SELECT COALESCE((extra->>'quota_used')::numeric, 0), extra ? 'quota_daily_used', extra ? 'quota_weekly_used' FROM accounts WHERE id = $1", + account.ID, + ).Scan("aUsed, &hasDailyUsed, &hasWeeklyUsed)) + require.InDelta(t, 3.75, quotaUsed, 0.000001) + require.False(t, hasDailyUsed) + require.False(t, hasWeeklyUsed) + + var outboxCount int + require.NoError(t, integrationDB.QueryRowContext(ctx, + "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2", + service.SchedulerOutboxEventAccountChanged, account.ID, + ).Scan(&outboxCount)) + require.Zero(t, outboxCount, "service_account under its limit must not enqueue scheduler outbox") +} + func TestUsageBillingRepositoryApply_EnqueuesSchedulerOutboxOnQuotaCrossing(t *testing.T) { ctx := context.Background() client := testEntClient(t) @@ -340,6 +444,140 @@ func TestUsageBillingRepositoryApply_EnqueuesSchedulerOutboxOnQuotaCrossing(t *t require.NoError(t, err) require.Equal(t, 1, outboxCountFor(t, accountID), "single-shot crossing weekly limit should enqueue once") }) + + t.Run("total_first_crossing_enqueues", func(t *testing.T) { + apiKeyID, accountID := newFixture(t, map[string]any{ + "quota_limit": 10.0, + }) + _, err := repo.Apply(ctx, &service.UsageBillingCommand{ + RequestID: uuid.NewString(), + APIKeyID: apiKeyID, + AccountID: accountID, + AccountType: service.AccountTypeAPIKey, + AccountQuotaCost: 4, + }) + require.NoError(t, err) + require.Equal(t, 0, outboxCountFor(t, accountID), "below limit should not enqueue") + + _, err = repo.Apply(ctx, &service.UsageBillingCommand{ + RequestID: uuid.NewString(), + APIKeyID: apiKeyID, + AccountID: accountID, + AccountType: service.AccountTypeAPIKey, + AccountQuotaCost: 8, + }) + require.NoError(t, err) + require.Equal(t, 1, outboxCountFor(t, accountID), "crossing total limit should enqueue once") + }) +} + +func TestIncrementUsageBillingAccountQuota_StatementVariantsMatchFullStatement(t *testing.T) { + ctx := context.Background() + client := testEntClient(t) + now := time.Now().UTC() + at := func(offset time.Duration) string { return now.Add(offset).Format(time.RFC3339) } + + // dailyUnset/weeklyUnset 为库中该维度视为未设置上限：省略该维度的语句在这些配置上命中，其余配置命中 0 行。 + configs := []struct { + name string + extra map[string]any + rawExtra string + dailyUnset bool + weeklyUnset bool + wantErr bool + }{ + {name: "no_limits", extra: map[string]any{}, dailyUnset: true, weeklyUnset: true}, + {name: "total_only", extra: map[string]any{"quota_limit": 100.0, "quota_used": 3.0}, dailyUnset: true, weeklyUnset: true}, + {name: "zero_and_negative_period_limits", extra: map[string]any{"quota_limit": 5.0, "quota_daily_limit": 0.0, "quota_weekly_limit": -1.0}, dailyUnset: true, weeklyUnset: true}, + {name: "null_period_limits", extra: map[string]any{"quota_daily_limit": nil, "quota_weekly_limit": nil}, dailyUnset: true, weeklyUnset: true}, + {name: "string_period_limits", extra: map[string]any{"quota_daily_limit": "10", "quota_weekly_limit": "0"}}, + {name: "daily_rolling_active", extra: map[string]any{"quota_daily_limit": 10.0, "quota_daily_used": 4.0, "quota_daily_start": at(-time.Hour)}, weeklyUnset: true}, + {name: "daily_rolling_expired", extra: map[string]any{"quota_daily_limit": 10.0, "quota_daily_used": 9.0, "quota_daily_start": at(-30 * time.Hour)}, weeklyUnset: true}, + {name: "daily_fixed_expired", extra: map[string]any{ + "quota_daily_limit": 10.0, "quota_daily_used": 7.0, "quota_daily_reset_mode": "fixed", + "quota_daily_reset_hour": 3, "quota_reset_timezone": "Asia/Shanghai", "quota_daily_reset_at": at(-2 * time.Hour), + }, weeklyUnset: true}, + {name: "weekly_rolling_active", extra: map[string]any{"quota_weekly_limit": 50.0, "quota_weekly_used": 20.0, "quota_weekly_start": at(-48 * time.Hour)}, dailyUnset: true}, + {name: "weekly_fixed_expired", extra: map[string]any{ + "quota_weekly_limit": 50.0, "quota_weekly_used": 49.0, "quota_weekly_reset_mode": "fixed", + "quota_weekly_reset_day": 1, "quota_weekly_reset_hour": 0, "quota_reset_timezone": "UTC", "quota_weekly_reset_at": at(-time.Hour), + }, dailyUnset: true}, + {name: "daily_and_weekly", extra: map[string]any{ + "quota_limit": 100.0, "quota_daily_limit": 10.0, "quota_weekly_limit": 50.0, + "quota_daily_used": 1.0, "quota_weekly_used": 2.0, "quota_daily_start": at(-time.Hour), "quota_weekly_start": at(-200 * time.Hour), + }}, + {name: "array_extra", rawExtra: `[]`}, + {name: "scalar_extra", rawExtra: `"value"`}, + {name: "non_numeric_daily_limit", extra: map[string]any{"quota_daily_limit": "abc"}, weeklyUnset: true, wantErr: true}, + {name: "boolean_weekly_limit", extra: map[string]any{"quota_weekly_limit": true}, dailyUnset: true, wantErr: true}, + } + + for _, cfg := range configs { + for _, hints := range [][2]bool{{false, false}, {true, false}, {false, true}, {true, true}} { + t.Run(fmt.Sprintf("%s/daily=%v/weekly=%v", cfg.name, hints[0], hints[1]), func(t *testing.T) { + variant := mustCreateAccount(t, client, &service.Account{ + Name: "usage-billing-quota-variant-" + uuid.NewString(), + Type: service.AccountTypeAPIKey, + Extra: maps.Clone(cfg.extra), + }) + reference := mustCreateAccount(t, client, &service.Account{ + Name: "usage-billing-quota-reference-" + uuid.NewString(), + Type: service.AccountTypeAPIKey, + Extra: maps.Clone(cfg.extra), + }) + if cfg.rawExtra != "" { + _, err := integrationDB.ExecContext(ctx, "UPDATE accounts SET extra = $1::jsonb WHERE id IN ($2, $3)", cfg.rawExtra, variant.ID, reference.ID) + require.NoError(t, err) + t.Cleanup(func() { + _, _ = integrationDB.ExecContext(ctx, "UPDATE accounts SET extra = '{}'::jsonb WHERE id IN ($1, $2)", variant.ID, reference.ID) + }) + } + + if cfg.wantErr { + sqlState := func(accountID int64, daily, weekly bool) string { + tx, err := integrationDB.BeginTx(ctx, nil) + require.NoError(t, err) + defer func() { _ = tx.Rollback() }() + _, err = incrementUsageBillingAccountQuota(ctx, tx, accountID, 2.5, daily, weekly) + var pgErr *pgconn.PgError + require.ErrorAs(t, err, &pgErr) + return pgErr.Code + } + require.Equal(t, sqlState(reference.ID, true, true), sqlState(variant.ID, hints[0], hints[1])) + return + } + + tx, err := integrationDB.BeginTx(ctx, nil) + require.NoError(t, err) + defer func() { _ = tx.Rollback() }() + + isObject := cfg.rawExtra == "" + wantHit := (hints[0] && hints[1]) || (isObject && (hints[0] || cfg.dailyUnset) && (hints[1] || cfg.weeklyUnset)) + _, err = tx.ExecContext(ctx, "SAVEPOINT quota_variant_probe") + require.NoError(t, err) + probe, err := execAccountQuotaIncrement(ctx, tx, accountQuotaIncrementSQLs[accountQuotaIncrementSQLIndex(hints[0], hints[1])], variant.ID, 2.5) + require.NoError(t, err) + require.Equal(t, wantHit, probe != nil, "variant statement hit") + _, err = tx.ExecContext(ctx, "ROLLBACK TO SAVEPOINT quota_variant_probe") + require.NoError(t, err) + + got, err := incrementUsageBillingAccountQuota(ctx, tx, variant.ID, 2.5, hints[0], hints[1]) + require.NoError(t, err) + want, err := incrementUsageBillingAccountQuota(ctx, tx, reference.ID, 2.5, true, true) + require.NoError(t, err) + require.Equal(t, *want, *got) + + var gotExtra, wantExtra string + var sameExtra, sameUpdatedAt bool + require.NoError(t, tx.QueryRowContext(ctx, + "SELECT v.extra::text, r.extra::text, v.extra = r.extra, v.updated_at = r.updated_at FROM accounts v, accounts r WHERE v.id = $1 AND r.id = $2", + variant.ID, reference.ID, + ).Scan(&gotExtra, &wantExtra, &sameExtra, &sameUpdatedAt)) + require.True(t, sameExtra, "variant extra %s != reference extra %s", gotExtra, wantExtra) + require.True(t, sameUpdatedAt) + }) + } + } } func TestDashboardAggregationRepositoryCleanupUsageBillingDedup_BatchDeletesOldRows(t *testing.T) { diff --git a/backend/internal/repository/usage_billing_repo_unit_test.go b/backend/internal/repository/usage_billing_repo_unit_test.go index 40394d819c8d..73f29d2e4cf0 100644 --- a/backend/internal/repository/usage_billing_repo_unit_test.go +++ b/backend/internal/repository/usage_billing_repo_unit_test.go @@ -5,6 +5,7 @@ package repository import ( "context" "database/sql" + "strings" "testing" "github.com/DATA-DOG/go-sqlmock" @@ -22,6 +23,7 @@ const ( userExistsForBillingSQL = `(?s)SELECT 1\s+FROM users\s+WHERE id = \$1 AND deleted_at IS NULL` apiKeyQuotaIncrementSQL = `(?s)UPDATE api_keys\s+SET quota_used = quota_used \+ \$1,.*WHERE id = \$2 AND deleted_at IS NULL\s+RETURNING` apiKeyRateLimitIncrementSQL = `(?s)UPDATE api_keys SET\s+usage_5h = .*WHERE id = \$2 AND deleted_at IS NULL` + accountQuotaIncrementSQL = `(?s)UPDATE accounts SET extra = .*WHERE id = \$2 AND deleted_at IS NULL.*RETURNING` ) func TestDeductUsageBillingBalance_UsesSufficientBalanceGuard(t *testing.T) { @@ -137,6 +139,42 @@ func TestApplyUsageBillingEffects_DeletedAPIKeyStillBillsBalance(t *testing.T) { require.NoError(t, mock.ExpectationsWereMet()) } +func TestApplyUsageBillingEffects_DeletedAccountStillBillsBalance(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(conditionalBalanceDeductSQL). + WithArgs(10.0, int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(90.0)) + mock.ExpectQuery(accountQuotaIncrementSQL). + WithArgs(4.0, int64(9)). + WillReturnRows(sqlmock.NewRows(accountQuotaStateColumns)) + mock.ExpectQuery(accountQuotaIncrementSQL). + WithArgs(4.0, int64(9)). + WillReturnRows(sqlmock.NewRows(accountQuotaStateColumns)) + mock.ExpectCommit() + + result := &service.UsageBillingApplyResult{Applied: true} + err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{ + UserID: 42, + AccountID: 9, + AccountType: service.AccountTypeServiceAccount, + BalanceCost: 10, + AccountQuotaCost: 4, + }, result) + require.NoError(t, err) + require.NotNil(t, result.NewBalance) + require.InDelta(t, 90.0, *result.NewBalance, 0.000001) + require.Nil(t, result.QuotaState) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} + func TestApplyUsageBillingEffects_APIKeyCounterErrorStillFails(t *testing.T) { ctx := context.Background() db, mock, err := sqlmock.New() @@ -184,6 +222,100 @@ func TestDeductUsageBillingBalance_ReturnsUserNotFoundWhenNoUserUpdated(t *testi require.NoError(t, mock.ExpectationsWereMet()) } +var accountQuotaStateColumns = []string{"quota_used", "quota_limit", "quota_daily_used", "quota_daily_limit", "quota_weekly_used", "quota_weekly_limit"} + +func TestAccountQuotaIncrementSQLs_GuardOmittedDimensions(t *testing.T) { + for _, daily := range []bool{false, true} { + for _, weekly := range []bool{false, true} { + query := accountQuotaIncrementSQLs[accountQuotaIncrementSQLIndex(daily, weekly)] + set, where, ok := strings.Cut(query, "WHERE ") + require.True(t, ok) + require.Equal(t, !daily || !weekly, strings.Contains(where, accountQuotaExtraIsObjectCond), "daily=%v weekly=%v", daily, weekly) + require.Equal(t, daily, strings.Contains(set, "'quota_daily_used'"), "daily=%v weekly=%v", daily, weekly) + require.Equal(t, weekly, strings.Contains(set, "'quota_weekly_used'"), "daily=%v weekly=%v", daily, weekly) + require.Equal(t, !daily, strings.Contains(where, accountQuotaDailyLimitUnsetCond), "daily=%v weekly=%v", daily, weekly) + require.Equal(t, !weekly, strings.Contains(where, accountQuotaWeeklyLimitUnsetCond), "daily=%v weekly=%v", daily, weekly) + } + } +} + +func TestIncrementUsageBillingAccountQuota_UsesStatementMatchingHints(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual)) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(accountQuotaIncrementSQLs[accountQuotaIncrementSQLIndex(true, false)]). + WithArgs(2.5, int64(9)). + WillReturnRows(sqlmock.NewRows(accountQuotaStateColumns).AddRow(12.5, 100.0, 2.5, 50.0, 0.0, 0.0)) + mock.ExpectCommit() + + state, err := incrementUsageBillingAccountQuota(ctx, tx, 9, 2.5, true, false) + require.NoError(t, err) + require.NotNil(t, state) + require.InDelta(t, 12.5, state.TotalUsed, 0.000001) + require.InDelta(t, 2.5, state.DailyUsed, 0.000001) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestIncrementUsageBillingAccountQuota_FallsBackToFullStatementWhenGuardMisses(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual)) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(accountQuotaIncrementSQLs[accountQuotaIncrementSQLIndex(false, false)]). + WithArgs(2.5, int64(9)). + WillReturnRows(sqlmock.NewRows(accountQuotaStateColumns)) + mock.ExpectQuery(accountQuotaIncrementSQLs[accountQuotaIncrementSQLIndex(true, true)]). + WithArgs(2.5, int64(9)). + WillReturnRows(sqlmock.NewRows(accountQuotaStateColumns).AddRow(2.5, 0.0, 0.0, 0.0, 2.5, 20.0)) + mock.ExpectCommit() + + state, err := incrementUsageBillingAccountQuota(ctx, tx, 9, 2.5, false, false) + require.NoError(t, err) + require.NotNil(t, state) + require.InDelta(t, 2.5, state.WeeklyUsed, 0.000001) + require.InDelta(t, 20.0, state.WeeklyLimit, 0.000001) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestIncrementUsageBillingAccountQuota_ReturnsNotFoundWhenFullStatementMisses(t *testing.T) { + for _, hints := range [][2]bool{{false, false}, {true, false}, {false, true}, {true, true}} { + ctx := context.Background() + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual)) + require.NoError(t, err) + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(accountQuotaIncrementSQLs[accountQuotaIncrementSQLIndex(hints[0], hints[1])]). + WithArgs(2.5, int64(9)). + WillReturnRows(sqlmock.NewRows(accountQuotaStateColumns)) + if !(hints[0] && hints[1]) { + mock.ExpectQuery(accountQuotaIncrementSQLs[accountQuotaIncrementSQLIndex(true, true)]). + WithArgs(2.5, int64(9)). + WillReturnRows(sqlmock.NewRows(accountQuotaStateColumns)) + } + mock.ExpectRollback() + + state, err := incrementUsageBillingAccountQuota(ctx, tx, 9, 2.5, hints[0], hints[1]) + require.ErrorIs(t, err, service.ErrAccountNotFound, "hints=%v", hints) + require.Nil(t, state) + require.NoError(t, tx.Rollback()) + require.NoError(t, mock.ExpectationsWereMet(), "hints=%v", hints) + _ = db.Close() + } +} + func TestReserveUsageBillingBatchImageBalance_MovesAvailableToFrozen(t *testing.T) { ctx := context.Background() db, mock, err := sqlmock.New() diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index f3e891bd0d1e..0972bdb07045 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -4,7 +4,6 @@ package service import ( "encoding/json" "errors" - "hash/fnv" "log/slog" "net/url" "reflect" @@ -83,6 +82,9 @@ type Account struct { modelMappingCacheRawLen int modelMappingCacheRawSig uint64 modelMappingCacheRuntimeVersion uint64 + // modelMappingCacheShared 为 true 时缓存来自只读共享母本（见 PrepareSharedSnapshot）， + // 指针、长度与运行时版本匹配即直接复用预解析结果，不计算签名。 + modelMappingCacheShared bool // header_overrides 热路径缓存（非持久化字段，同 model_mapping 缓存先例） headerOverrideCache map[string]string @@ -238,7 +240,7 @@ func (a *Account) IsSchedulable() bool { if a.TempUnschedulableUntil != nil && now.Before(*a.TempUnschedulableUntil) { return false } - if a.IsAPIKeyOrBedrock() && a.IsQuotaExceeded() { + if a.SupportsAccountQuota() && a.IsQuotaExceeded() { return false } return true @@ -667,6 +669,9 @@ func (a *Account) GetModelMapping() map[string]string { a.modelMappingCacheRawPtr == rawPtr && a.modelMappingCacheRawLen == rawLen && a.modelMappingCacheRuntimeVersion == runtimeVersion { + if a.modelMappingCacheShared { + return a.modelMappingCache + } rawSig = modelMappingSignature(rawMapping) rawSigReady = true if a.modelMappingCacheRawSig == rawSig { @@ -686,9 +691,24 @@ func (a *Account) GetModelMapping() map[string]string { a.modelMappingCacheRawLen = rawLen a.modelMappingCacheRawSig = rawSig a.modelMappingCacheRuntimeVersion = runtimeVersion + a.modelMappingCacheShared = false return mapping } +// PrepareSharedSnapshot 预解析模型映射，并把账号标记为只读共享母本。 +// 母本按值复制出的副本与母本共用 Credentials 及其中 model_mapping 的同一 map： +// 这些内层 map 在运行期不被原地修改，母本留在缓存期间持有这些 map，其地址不会被复用， +// 因此副本在指针、长度与运行时版本都匹配时直接复用预解析结果，不计算签名； +// 副本替换 Credentials 等导致任一项不匹配时回到普通的重建路径。 +// 只能由构造该账号的 goroutine 在其发布给其它 goroutine 之前调用，发布后母本不得再被修改。 +func (a *Account) PrepareSharedSnapshot() { + if a == nil { + return + } + a.GetModelMapping() + a.modelMappingCacheShared = true +} + func (a *Account) resolveModelMapping(rawMapping map[string]any) map[string]string { if a.Credentials == nil { // Antigravity 平台使用默认映射 @@ -749,28 +769,56 @@ func mapPtr(m map[string]any) uintptr { return reflect.ValueOf(m).Pointer() } +const ( + modelMappingSignatureFNVOffset = 14695981039346656037 + modelMappingSignatureFNVPrime = 1099511628211 +) + +// modelMappingSignature 计算原始映射内容的进程内校验值，与 map 遍历顺序无关且不分配内存。 +// 每个条目把 key、分隔字节 0 与 value（string 写其字节，其它类型写标记字节 1） +// 放进同一个 FNV-1a 哈希并经 fmix64 收尾，再对所有条目求和； +// key 与 value 同在一个条目哈希里，交换两个 key 的 value 会改变结果。 func modelMappingSignature(rawMapping map[string]any) uint64 { if len(rawMapping) == 0 { return 0 } - keys := make([]string, 0, len(rawMapping)) - for k := range rawMapping { - keys = append(keys, k) - } - sort.Strings(keys) - - h := fnv.New64a() - for _, k := range keys { - _, _ = h.Write([]byte(k)) - _, _ = h.Write([]byte{0}) - if v, ok := rawMapping[k].(string); ok { - _, _ = h.Write([]byte(v)) + var sum uint64 + for k, v := range rawMapping { + h := uint64(modelMappingSignatureFNVOffset) + h = fnv1aString(h, k) + h = fnv1aByte(h, 0) + if s, ok := v.(string); ok { + h = fnv1aString(h, s) } else { - _, _ = h.Write([]byte{1}) + h = fnv1aByte(h, 1) } - _, _ = h.Write([]byte{0xff}) + sum += fmix64(h) } - return h.Sum64() + return sum +} + +func fnv1aString(h uint64, s string) uint64 { + for i := 0; i < len(s); i++ { + h ^= uint64(s[i]) + h *= modelMappingSignatureFNVPrime + } + return h +} + +func fnv1aByte(h uint64, b byte) uint64 { + h ^= uint64(b) + h *= modelMappingSignatureFNVPrime + return h +} + +// fmix64 是 MurmurHash3 的 64 位收尾函数，使条目哈希各位充分混合后再求和。 +func fmix64(h uint64) uint64 { + h ^= h >> 33 + h *= 0xff51afd7ed558ccd + h ^= h >> 33 + h *= 0xc4ceb9fe1a85ec53 + h ^= h >> 33 + return h } // antigravityDefaultPassthroughModels 是 Antigravity 显式 model_mapping 运行时补齐的默认透传条目： @@ -1410,11 +1458,26 @@ func (a *Account) IsBedrockAPIKey() bool { return a.IsBedrock() && a.GetCredential("auth_mode") == "apikey" } -// IsAPIKeyOrBedrock 返回账号类型是否支持配额和池模式等特性 +// IsAPIKeyOrBedrock 返回账号类型是否支持池模式等特性（账号配额判定见 SupportsAccountQuota） func (a *Account) IsAPIKeyOrBedrock() bool { return a.Type == AccountTypeAPIKey || a.Type == AccountTypeBedrock } +// AccountTypeSupportsQuota 返回账号类型是否支持账号配额（quota_limit 及日/周额度、quota_used 累加、超限停调度）。 +// 在 apikey/bedrock 之外额外放开 Vertex service_account；池模式等其它特性仍只看 IsAPIKeyOrBedrock。 +func AccountTypeSupportsQuota(accountType string) bool { + switch strings.ToLower(accountType) { + case AccountTypeAPIKey, AccountTypeBedrock, AccountTypeServiceAccount: + return true + } + return false +} + +// SupportsAccountQuota 返回账号是否支持账号配额，类型集合见 AccountTypeSupportsQuota。 +func (a *Account) SupportsAccountQuota() bool { + return AccountTypeSupportsQuota(a.Type) +} + func (a *Account) IsOpenAI() bool { return a != nil && a.Platform == PlatformOpenAI } diff --git a/backend/internal/service/account_model_mapping_signature_test.go b/backend/internal/service/account_model_mapping_signature_test.go new file mode 100644 index 000000000000..a83c05dc81ab --- /dev/null +++ b/backend/internal/service/account_model_mapping_signature_test.go @@ -0,0 +1,96 @@ +//go:build unit + +package service + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestModelMappingSignature_IndependentOfInsertionOrder(t *testing.T) { + const size = 64 + forward := make(map[string]any, size) + backward := make(map[string]any, size) + for i := 0; i < size; i++ { + forward[fmt.Sprintf("model-%02d", i)] = fmt.Sprintf("upstream-%02d", i) + } + for i := size - 1; i >= 0; i-- { + backward[fmt.Sprintf("model-%02d", i)] = fmt.Sprintf("upstream-%02d", i) + } + forward["non-string"] = 42 + backward["non-string"] = 42 + + want := modelMappingSignature(forward) + require.NotZero(t, want) + for i := 0; i < 20; i++ { + require.Equal(t, want, modelMappingSignature(forward)) + require.Equal(t, want, modelMappingSignature(backward)) + } +} + +func TestModelMappingSignature_DetectsContentChanges(t *testing.T) { + base := map[string]any{ + "claude-sonnet": "sonnet-a", + "claude-opus": "opus-a", + "gpt-5": "gpt-5-upstream", + } + baseSig := modelMappingSignature(base) + + cases := map[string]map[string]any{ + "value changed": { + "claude-sonnet": "sonnet-b", + "claude-opus": "opus-a", + "gpt-5": "gpt-5-upstream", + }, + "key renamed": { + "claude-sonnet-x": "sonnet-a", + "claude-opus": "opus-a", + "gpt-5": "gpt-5-upstream", + }, + "values swapped between keys": { + "claude-sonnet": "opus-a", + "claude-opus": "sonnet-a", + "gpt-5": "gpt-5-upstream", + }, + "string value replaced by non-string": { + "claude-sonnet": 1, + "claude-opus": "opus-a", + "gpt-5": "gpt-5-upstream", + }, + "entry removed": { + "claude-sonnet": "sonnet-a", + "claude-opus": "opus-a", + }, + } + for name, changed := range cases { + require.NotEqual(t, baseSig, modelMappingSignature(changed), name) + } + + // 非 string 的 value 统一按类型标记参与签名，具体取值不影响结果。 + require.Equal(t, + modelMappingSignature(map[string]any{"m": 1}), + modelMappingSignature(map[string]any{"m": true}), + ) +} + +func TestModelMappingSignature_EmptyIsZero(t *testing.T) { + require.Zero(t, modelMappingSignature(nil)) + require.Zero(t, modelMappingSignature(map[string]any{})) +} + +func TestModelMappingSignature_DoesNotAllocate(t *testing.T) { + mapping := map[string]any{ + "claude-sonnet": "sonnet-a", + "claude-opus": "opus-a", + "gpt-5": "gpt-5-upstream", + "non-string": 3.14, + } + var sink uint64 + allocs := testing.AllocsPerRun(1000, func() { + sink += modelMappingSignature(mapping) + }) + require.Zero(t, allocs) + require.NotZero(t, sink) +} diff --git a/backend/internal/service/account_quota_schedulable_test.go b/backend/internal/service/account_quota_schedulable_test.go index 2895b34c889e..dfe34c4fc507 100644 --- a/backend/internal/service/account_quota_schedulable_test.go +++ b/backend/internal/service/account_quota_schedulable_test.go @@ -113,6 +113,48 @@ func TestAccountIsSchedulable_QuotaExceeded(t *testing.T) { }, want: false, }, + { + name: "service_account total quota exceeded", + account: &Account{ + Status: StatusActive, + Schedulable: true, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Extra: map[string]any{ + "quota_limit": 50.0, + "quota_used": 50.0, + }, + }, + want: false, + }, + { + name: "service_account daily quota exceeded", + account: &Account{ + Status: StatusActive, + Schedulable: true, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Extra: map[string]any{ + "quota_daily_limit": 10.0, + "quota_daily_used": 12.0, + "quota_daily_start": now.Add(-1 * time.Hour).Format(time.RFC3339), + }, + }, + want: false, + }, + { + name: "service_account without limit stays schedulable regardless of quota_used", + account: &Account{ + Status: StatusActive, + Schedulable: true, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Extra: map[string]any{ + "quota_used": 9999.0, + }, + }, + want: true, + }, } for _, tt := range tests { diff --git a/backend/internal/service/account_quota_service_account_test.go b/backend/internal/service/account_quota_service_account_test.go new file mode 100644 index 000000000000..6358f3202417 --- /dev/null +++ b/backend/internal/service/account_quota_service_account_test.go @@ -0,0 +1,204 @@ +package service + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAccountSupportsAccountQuota(t *testing.T) { + tests := []struct { + accountType string + want bool + }{ + {AccountTypeAPIKey, true}, + {AccountTypeBedrock, true}, + {AccountTypeServiceAccount, true}, + {AccountTypeOAuth, false}, + {AccountTypeSetupToken, false}, + {AccountTypeUpstream, false}, + } + for _, tt := range tests { + t.Run(tt.accountType, func(t *testing.T) { + account := &Account{Type: tt.accountType} + require.Equal(t, tt.want, account.SupportsAccountQuota()) + }) + } +} + +// 配额专用判定放开 service_account，不能连带放开池模式等其它 IsAPIKeyOrBedrock 特性。 +func TestServiceAccountQuotaDoesNotWidenPoolMode(t *testing.T) { + account := &Account{ + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: map[string]any{"pool_mode": true}, + } + require.True(t, account.SupportsAccountQuota()) + require.False(t, account.IsAPIKeyOrBedrock()) + require.False(t, account.IsPoolMode()) +} + +func TestShouldUpdateAccountQuota(t *testing.T) { + tests := []struct { + name string + account *Account + totalCost float64 + want bool + }{ + { + name: "apikey without limit does not accumulate", + account: &Account{Type: AccountTypeAPIKey, Extra: map[string]any{}}, + totalCost: 1.5, + want: false, + }, + { + name: "apikey with total limit accumulates", + account: &Account{Type: AccountTypeAPIKey, Extra: map[string]any{"quota_limit": 100.0}}, + totalCost: 1.5, + want: true, + }, + { + name: "bedrock with daily limit accumulates", + account: &Account{Type: AccountTypeBedrock, Extra: map[string]any{"quota_daily_limit": 10.0}}, + totalCost: 1.5, + want: true, + }, + { + name: "service_account without limit does not accumulate", + account: &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Extra: map[string]any{"quota_used": 30.0}}, + totalCost: 1.5, + want: false, + }, + { + name: "service_account with limit accumulates", + account: &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Extra: map[string]any{"quota_weekly_limit": 50.0}}, + totalCost: 1.5, + want: true, + }, + { + name: "service_account zero cost does not accumulate", + account: &Account{Platform: PlatformGemini, Type: AccountTypeServiceAccount, Extra: map[string]any{"quota_limit": 100.0}}, + totalCost: 0, + want: false, + }, + { + name: "oauth with limit is not quota eligible", + account: &Account{Type: AccountTypeOAuth, Extra: map[string]any{"quota_limit": 100.0}}, + totalCost: 1.5, + want: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + p := &postUsageBillingParams{ + Cost: &CostBreakdown{TotalCost: tt.totalCost}, + Account: tt.account, + } + require.Equal(t, tt.want, p.shouldUpdateAccountQuota()) + }) + } +} + +func TestBuildUsageBillingCommandAccountQuotaCostForServiceAccount(t *testing.T) { + build := func(account *Account) *UsageBillingCommand { + return buildUsageBillingCommand("req-sa-quota", nil, &postUsageBillingParams{ + Cost: &CostBreakdown{TotalCost: 2, ActualCost: 3}, + User: &User{ID: 7}, + APIKey: &APIKey{ID: 13}, + Account: account, + AccountRateMultiplier: 1.5, + }) + } + + // service_account 设了上限：AccountQuotaCost = TotalCost × 账号倍率 + cmd := build(&Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeServiceAccount, Extra: map[string]any{"quota_limit": 100.0}}) + require.NotNil(t, cmd) + require.Equal(t, AccountTypeServiceAccount, cmd.AccountType) + require.InDelta(t, 3.0, cmd.AccountQuotaCost, 1e-9) + + // service_account 不设上限：与 apikey 一致不累加 + cmd = build(&Account{ID: 10, Platform: PlatformGemini, Type: AccountTypeServiceAccount}) + require.NotNil(t, cmd) + require.Zero(t, cmd.AccountQuotaCost) +} + +type accountQuotaIncrementRepoStub struct { + AccountRepository + calls []float64 +} + +func (s *accountQuotaIncrementRepoStub) IncrementQuotaUsed(_ context.Context, _ int64, amount float64) error { + s.calls = append(s.calls, amount) + return nil +} + +func TestPostUsageBillingLegacyAccumulatesServiceAccountQuota(t *testing.T) { + repo := &accountQuotaIncrementRepoStub{} + postUsageBilling(context.Background(), &postUsageBillingParams{ + Cost: &CostBreakdown{TotalCost: 2}, + User: &User{ID: 7}, + APIKey: &APIKey{ID: 13}, + Account: &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeServiceAccount, Extra: map[string]any{"quota_limit": 100.0}}, + AccountRateMultiplier: 1.25, + }, &billingDeps{accountRepo: repo}) + + require.Len(t, repo.calls, 1) + require.InDelta(t, 2.5, repo.calls[0], 1e-9) +} + +type settingGetValueCounter struct { + SettingRepository + calls int +} + +func (s *settingGetValueCounter) GetValue(context.Context, string) (string, error) { + s.calls++ + return "false", nil +} + +// 未设任何上限的账号没有可跨越的阈值，不应逐请求读取通知设置。 +func TestNotifyAccountQuotaSkipsAccountWithoutLimit(t *testing.T) { + tests := []struct { + name string + account *Account + wantCalls int + }{ + { + name: "service_account without limit", + account: &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeServiceAccount}, + wantCalls: 0, + }, + { + name: "service_account with total limit", + account: &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeServiceAccount, Extra: map[string]any{"quota_limit": 100.0}}, + wantCalls: 1, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + settings := &settingGetValueCounter{} + notifyAccountQuota(&postUsageBillingParams{ + Cost: &CostBreakdown{TotalCost: 2}, + Account: tt.account, + AccountRateMultiplier: 1, + }, &billingDeps{ + balanceNotifyService: NewBalanceNotifyService(NewEmailService(settings, nil), settings, nil), + }, nil) + require.Equal(t, tt.wantCalls, settings.calls) + }) + } +} + +// 账号已删除时计数被跳过、无递增后状态，不应回退到快照做阈值判定。 +func TestNotifyAccountQuotaSkipsWhenQuotaIncrementSkipped(t *testing.T) { + settings := &settingGetValueCounter{} + notifyAccountQuota(&postUsageBillingParams{ + Cost: &CostBreakdown{TotalCost: 2}, + Account: &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeServiceAccount, Extra: map[string]any{"quota_limit": 100.0}}, + AccountRateMultiplier: 1, + }, &billingDeps{ + balanceNotifyService: NewBalanceNotifyService(NewEmailService(settings, nil), settings, nil), + }, &UsageBillingApplyResult{Applied: true}) + require.Zero(t, settings.calls) +} diff --git a/backend/internal/service/account_shared_snapshot_test.go b/backend/internal/service/account_shared_snapshot_test.go new file mode 100644 index 000000000000..2d3421193d6d --- /dev/null +++ b/backend/internal/service/account_shared_snapshot_test.go @@ -0,0 +1,82 @@ +//go:build unit + +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// 共享母本的副本命中预解析缓存时不计算签名：此处故意原地改写共享的原始映射， +// 只用来观察副本走的是免签名路径（生产代码不会原地修改该 map）。 +func TestAccountPrepareSharedSnapshot_CopiesReusePreparedMappingWithoutSignature(t *testing.T) { + rawMapping := map[string]any{"claude-sonnet": "sonnet-a"} + master := &Account{ + Platform: PlatformAnthropic, + Credentials: map[string]any{"model_mapping": rawMapping}, + } + master.PrepareSharedSnapshot() + + replica := *master + require.Equal(t, map[string]string{"claude-sonnet": "sonnet-a"}, replica.GetModelMapping()) + + rawMapping["claude-sonnet"] = "sonnet-b" + require.Equal(t, "sonnet-a", replica.GetModelMapping()["claude-sonnet"], + "the shared fast path must not recompute the signature") + require.Equal(t, "sonnet-a", replica.GetMappedModel("claude-sonnet")) + + // 未标记为共享的账号仍按签名检出原地改写。 + plainRaw := map[string]any{"claude-sonnet": "sonnet-a"} + plain := &Account{ + Platform: PlatformAnthropic, + Credentials: map[string]any{"model_mapping": plainRaw}, + } + require.Equal(t, "sonnet-a", plain.GetModelMapping()["claude-sonnet"]) + plainRaw["claude-sonnet"] = "sonnet-b" + require.Equal(t, "sonnet-b", plain.GetModelMapping()["claude-sonnet"]) +} + +func TestAccountPrepareSharedSnapshot_ReplacedCredentialsFallBackToSignaturePath(t *testing.T) { + master := &Account{ + Platform: PlatformAnthropic, + Credentials: map[string]any{"model_mapping": map[string]any{"claude-sonnet": "sonnet-a"}}, + } + master.PrepareSharedSnapshot() + + replica := *master + replacedRaw := map[string]any{"claude-sonnet": "sonnet-b"} + replica.Credentials = map[string]any{"model_mapping": replacedRaw} + require.Equal(t, "sonnet-b", replica.GetModelMapping()["claude-sonnet"]) + + // 重建后副本回到普通路径：同一 map 的原地改写由签名检出。 + replacedRaw["claude-sonnet"] = "sonnet-c" + require.Equal(t, "sonnet-c", replica.GetModelMapping()["claude-sonnet"]) + + // 母本不受副本重建影响。 + require.Equal(t, "sonnet-a", master.GetModelMapping()["claude-sonnet"]) +} + +func TestAccountPrepareSharedSnapshot_MappingLengthChangeRebuilds(t *testing.T) { + rawMapping := map[string]any{"claude-sonnet": "sonnet-a"} + master := &Account{ + Platform: PlatformAnthropic, + Credentials: map[string]any{"model_mapping": rawMapping}, + } + master.PrepareSharedSnapshot() + + replica := *master + rawMapping["claude-opus"] = "opus-a" + require.Equal(t, "opus-a", replica.GetModelMapping()["claude-opus"]) +} + +func TestAccountPrepareSharedSnapshot_NilAccountAndNoMapping(t *testing.T) { + var nilAccount *Account + require.NotPanics(t, nilAccount.PrepareSharedSnapshot) + + master := &Account{Platform: PlatformAnthropic} + master.PrepareSharedSnapshot() + replica := *master + require.Nil(t, replica.GetModelMapping()) + require.Equal(t, "claude-sonnet", replica.GetMappedModel("claude-sonnet")) +} diff --git a/backend/internal/service/account_test_service_openai_image_test.go b/backend/internal/service/account_test_service_openai_image_test.go index 8ad5c20fb058..f444d767eecd 100644 --- a/backend/internal/service/account_test_service_openai_image_test.go +++ b/backend/internal/service/account_test_service_openai_image_test.go @@ -61,7 +61,7 @@ func TestAccountTestService_OpenAIImageOAuthHandlesOutputItemDoneFallback(t *tes require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.lastReq.Context())) // Force 优先于自定义 UA；无种子账号使用规范身份。压缩后仍是 image_generation 工具调用。 require.Equal(t, resolveCodexOutboundIdentityForAccount(account, "").userAgent, upstream.lastReq.Header.Get("User-Agent")) - require.Equal(t, "codex-tui", upstream.lastReq.Header.Get("originator")) + require.Equal(t, "codex_cli_rs", upstream.lastReq.Header.Get("originator")) require.Equal(t, "zstd", upstream.lastReq.Header.Get("Content-Encoding")) require.Equal(t, "image_generation", gjson.GetBytes(zstdDecodeForTest(t, upstream.lastBody), "tools.0.type").String()) require.Contains(t, rec.Body.String(), "Calling Codex /responses image tool") diff --git a/backend/internal/service/admin_group.go b/backend/internal/service/admin_group.go index 96c8092a8ec0..f2c06a9a839a 100644 --- a/backend/internal/service/admin_group.go +++ b/backend/internal/service/admin_group.go @@ -311,10 +311,17 @@ func defaultModelsListCandidateIDs(platform string) []string { } } -func defaultAllowImageGenerationForPlatform(platform string) bool { - // Grok image and video generation routes share the legacy image-generation gate. - // Older clients send the false zero value, so Grok groups must default enabled. - return platform == PlatformGrok +// DefaultAllowImageGenerationForPlatform reports whether new groups on the platform allow +// image generation. CreateGroup ORs the requested value with it, so a false in the create +// request does not turn it off. Composite defaults disabled because the same switch also +// gates OpenAI image generation there. +func DefaultAllowImageGenerationForPlatform(platform string) bool { + switch platform { + case PlatformGrok, PlatformGemini, PlatformAntigravity: + return true + default: + return false + } } func compositeDefaultModelsListCandidateIDs() []string { @@ -525,7 +532,7 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn mcpXMLInject = *input.MCPXMLInject } - allowImageGeneration := input.AllowImageGeneration || defaultAllowImageGenerationForPlatform(platform) + allowImageGeneration := input.AllowImageGeneration || DefaultAllowImageGenerationForPlatform(platform) allowBatchImageGeneration := input.AllowBatchImageGeneration && allowImageGeneration && platform == PlatformGemini // 如果指定了复制账号的源分组，先获取账号 ID 列表 diff --git a/backend/internal/service/admin_service_group_test.go b/backend/internal/service/admin_service_group_test.go index d87e130205fc..6ad7000f20f4 100644 --- a/backend/internal/service/admin_service_group_test.go +++ b/backend/internal/service/admin_service_group_test.go @@ -596,41 +596,47 @@ func TestAdminService_CreateGroup_DefaultsGrokMediaGenerationEnabled(t *testing. require.True(t, group.AllowImageGeneration) } -func TestAdminService_CreateGroup_PreservesNonGrokImageGenerationDisabled(t *testing.T) { - repo := &groupRepoStubForAdmin{} - svc := &adminServiceImpl{groupRepo: repo} +func TestAdminService_CreateGroup_DefaultsGeminiFamilyImageGenerationEnabled(t *testing.T) { + for _, platform := range []string{PlatformGemini, PlatformAntigravity} { + t.Run(platform, func(t *testing.T) { + repo := &groupRepoStubForAdmin{} + svc := &adminServiceImpl{groupRepo: repo} - group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{ - Name: "anthropic-text", - Description: "Anthropic text group", - Platform: PlatformAnthropic, - RateMultiplier: 1.0, - }) - require.NoError(t, err) - require.NotNil(t, group) - require.NotNil(t, repo.created) - require.False(t, repo.created.AllowImageGeneration) - require.False(t, group.AllowImageGeneration) + group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{ + Name: platform + "-images", + Description: "Gemini family group", + Platform: platform, + RateMultiplier: 1.0, + }) + require.NoError(t, err) + require.NotNil(t, group) + require.NotNil(t, repo.created) + require.True(t, repo.created.AllowImageGeneration) + require.True(t, group.AllowImageGeneration) + require.False(t, repo.created.AllowBatchImageGeneration) + }) + } } -func TestAdminService_CreateGroup_DisablesBatchImageWhenImageGenerationDisabled(t *testing.T) { - repo := &groupRepoStubForAdmin{} - svc := &adminServiceImpl{groupRepo: repo} +func TestAdminService_CreateGroup_PreservesImageGenerationDisabledWithoutPlatformDefault(t *testing.T) { + for _, platform := range []string{PlatformAnthropic, PlatformComposite} { + t.Run(platform, func(t *testing.T) { + repo := &groupRepoStubForAdmin{} + svc := &adminServiceImpl{groupRepo: repo} - group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{ - Name: "gemini-no-image", - Description: "Gemini group without image generation", - Platform: PlatformGemini, - RateMultiplier: 1.0, - AllowImageGeneration: false, - AllowBatchImageGeneration: true, - }) - require.NoError(t, err) - require.NotNil(t, group) - require.NotNil(t, repo.created) - require.False(t, repo.created.AllowImageGeneration) - require.False(t, repo.created.AllowBatchImageGeneration) - require.False(t, group.AllowBatchImageGeneration) + group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{ + Name: platform + "-text", + Description: "Group without image default", + Platform: platform, + RateMultiplier: 1.0, + }) + require.NoError(t, err) + require.NotNil(t, group) + require.NotNil(t, repo.created) + require.False(t, repo.created.AllowImageGeneration) + require.False(t, group.AllowImageGeneration) + }) + } } func TestAdminService_CreateGroup_DisablesBatchImageForNonGeminiPlatform(t *testing.T) { diff --git a/backend/internal/service/agent_upstream_error_scope_test.go b/backend/internal/service/agent_upstream_error_scope_test.go index 7cbbeabc310c..e79bdcd06be4 100644 --- a/backend/internal/service/agent_upstream_error_scope_test.go +++ b/backend/internal/service/agent_upstream_error_scope_test.go @@ -10,24 +10,6 @@ import ( "github.com/stretchr/testify/require" ) -func TestAgentReportedUpstream400IsNotClientValidation(t *testing.T) { - s := &OpenAIGatewayService{} - for _, tc := range []struct { - name, body string - retry bool - }{ - {"nested_upstream_failure", `{"error":{"code":"400","message":"400","param":"","type":"upstream_error"}}`, true}, - {"client_validation", `{"error":{"code":"invalid_request_error","message":"invalid input","type":"invalid_request_error"}}`, false}, - {"parameter_identified", `{"error":{"code":"400","param":"tools","type":"upstream_error"}}`, false}, - {"opaque_number_only", `{"error":{"code":"400","message":"400"}}`, false}, - {"continuation", `{"error":{"code":"invalid_encrypted_content","type":"upstream_error"}}`, false}, - } { - t.Run(tc.name, func(t *testing.T) { - require.Equal(t, tc.retry, s.shouldFailoverOpenAIUpstreamResponse(nil, 400, "", []byte(tc.body))) - }) - } -} - func TestAgentRequestBudgetRejectionDoesNotBecomeAccountAuthCooldown(t *testing.T) { a := &Account{ID: 7, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Credentials: map[string]any{"pool_mode": true, "base_url": "https://fixture.example"}} s := &OpenAIGatewayService{} diff --git a/backend/internal/service/antigravity_gateway_claude.go b/backend/internal/service/antigravity_gateway_claude.go index 264cbe998094..70230dd48888 100644 --- a/backend/internal/service/antigravity_gateway_claude.go +++ b/backend/internal/service/antigravity_gateway_claude.go @@ -61,6 +61,9 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, // 应用 thinking 模式自动后缀：如果 thinking 开启且目标是 claude-sonnet-4-5，自动改为 thinking 版本 thinkingEnabled := claudeReq.Thinking != nil && (claudeReq.Thinking.Type == "enabled" || claudeReq.Thinking.Type == "adaptive") mappedModel = applyThinkingModelSuffix(mappedModel, thinkingEnabled) + if geminiImageGenerationDenied(c, nil, originalModel, mappedModel) { + return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", ImageGenerationPermissionMessage()) + } billingModel := mappedModel // 获取 access_token diff --git a/backend/internal/service/antigravity_gateway_compat.go b/backend/internal/service/antigravity_gateway_compat.go index 60ecd977ed5e..2869675449f7 100644 --- a/backend/internal/service/antigravity_gateway_compat.go +++ b/backend/internal/service/antigravity_gateway_compat.go @@ -225,6 +225,9 @@ func (s *AntigravityGatewayService) prepareAntigravityCompatCall( thinkingEnabled := claudeRequest.Thinking != nil && (claudeRequest.Thinking.Type == "enabled" || claudeRequest.Thinking.Type == "adaptive") mappedModel = applyThinkingModelSuffix(mappedModel, thinkingEnabled) + if geminiImageGenerationDenied(c, nil, request.originalModel, mappedModel) { + return nil, s.writeAntigravityCompatError(c, http.StatusBadRequest, "invalid_request_error", ImageGenerationPermissionMessage()) + } if s.tokenProvider == nil { return nil, s.writeAntigravityCompatError(c, http.StatusBadGateway, "api_error", "Antigravity token provider not configured") diff --git a/backend/internal/service/antigravity_gateway_compat_stream.go b/backend/internal/service/antigravity_gateway_compat_stream.go index 217aaa17c2fe..140792451d63 100644 --- a/backend/internal/service/antigravity_gateway_compat_stream.go +++ b/backend/internal/service/antigravity_gateway_compat_stream.go @@ -20,14 +20,11 @@ type antigravityCompatStreamAdapter interface { Emit(*apicompat.AnthropicStreamEvent, *antigravityClientWriter) Finalize(*antigravityClientWriter) WriteError(*antigravityClientWriter, string) - Err() error } type antigravityChatStreamAdapter struct { anthropicState *apicompat.AnthropicEventToResponsesState chatState *apicompat.ResponsesEventToChatState - err error - done bool } func newAntigravityChatStreamAdapter(model string, includeUsage bool) *antigravityChatStreamAdapter { @@ -43,55 +40,29 @@ func newAntigravityChatStreamAdapter(model string, includeUsage bool) *antigravi } func (a *antigravityChatStreamAdapter) Emit(event *apicompat.AnthropicStreamEvent, writer *antigravityClientWriter) { - if a.done { - return - } for _, responseEvent := range apicompat.AnthropicEventToResponsesEvents(event, a.anthropicState) { a.emitResponseEvent(&responseEvent, writer) } } func (a *antigravityChatStreamAdapter) Finalize(writer *antigravityClientWriter) { - if a.done { - return - } for _, responseEvent := range apicompat.FinalizeAnthropicResponsesStream(a.anthropicState) { a.emitResponseEvent(&responseEvent, writer) } - if a.err != nil { - return - } for _, chunk := range apicompat.FinalizeResponsesChatStream(a.chatState) { if data, err := apicompat.ChatChunkToSSE(chunk); err == nil { writer.Write([]byte(data)) } } writer.Write([]byte("data: [DONE]\n\n")) - a.done = true } func (a *antigravityChatStreamAdapter) WriteError(writer *antigravityClientWriter, reason string) { - if a.done { - return - } writer.Fprintf("data: {\"error\":{\"message\":%q,\"type\":\"upstream_error\"}}\n\n", reason) - a.done = true } -func (a *antigravityChatStreamAdapter) Err() error { return a.err } - func (a *antigravityChatStreamAdapter) emitResponseEvent(event *apicompat.ResponsesStreamEvent, writer *antigravityClientWriter) { - if a.done { - return - } - chunks := apicompat.ResponsesEventToChatChunks(event, a.chatState) - if a.chatState.ProtocolError != "" { - a.err = errors.New(a.chatState.ProtocolError) - a.done = true - writer.Write([]byte(buildChatStreamErrorSSE("upstream_protocol_error", a.chatState.ProtocolError) + "data: [DONE]\n\n")) - return - } - for _, chunk := range chunks { + for _, chunk := range apicompat.ResponsesEventToChatChunks(event, a.chatState) { if data, err := apicompat.ChatChunkToSSE(chunk); err == nil { writer.Write([]byte(data)) } @@ -124,8 +95,6 @@ func (a *antigravityResponsesStreamAdapter) WriteError(writer *antigravityClient writer.Fprintf("event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"upstream_error\",\"message\":%q}}\n\n", reason) } -func (a *antigravityResponsesStreamAdapter) Err() error { return nil } - func (a *antigravityResponsesStreamAdapter) emitResponseEvent(event apicompat.ResponsesStreamEvent, writer *antigravityClientWriter) { if data, err := apicompat.ResponsesEventToSSE(event); err == nil { writer.Write([]byte(data)) @@ -378,12 +347,7 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStreamWithKeepaliveIn if !session.hasMeaningfulData() && !writer.Disconnected() { return handleAntigravityCompatEmptyStream(c, session) } - result, finishErr := session.finish() - if err := adapter.Err(); err != nil { - MarkResponseCommitted(c) - return result, err - } - return result, finishErr + return session.finish() } if event.err != nil { return s.handleAntigravityCompatReadError(c, session, event.err, maxLineSize, prefix) @@ -391,14 +355,8 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStreamWithKeepaliveIn resetAntigravityCompatTimer(timeoutTimer, timeout) s.observeAntigravityGeminiSSELine(c, event.line) session.consume(event.line) - if adapter.Err() != nil { - MarkResponseCommitted(c) - } case <-timeoutCh: - if err := adapter.Err(); err != nil { - return session.collectResult(writer.Disconnected()), err - } if writer.Disconnected() { return session.collectResult(true), nil } @@ -410,7 +368,7 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStreamWithKeepaliveIn return session.collectResult(false), fmt.Errorf("stream data interval timeout") case <-keepaliveCh: - if session.hasMeaningfulData() && !writer.Disconnected() && adapter.Err() == nil { + if session.hasMeaningfulData() && !writer.Disconnected() { writer.Write([]byte(": ping\n\n")) } case now := <-preContentTimer.C: @@ -520,9 +478,6 @@ func (s *AntigravityGatewayService) handleAntigravityCompatReadError( maxLineSize int, prefix string, ) (*antigravityStreamResult, error) { - if protocolErr := session.adapter.Err(); protocolErr != nil { - return session.collectResult(session.writer.Disconnected()), protocolErr - } if !session.hasMeaningfulData() && !session.writer.Disconnected() { if session.preContentKeepaliveSent { writeAntigravityCompatStreamError(c, session.adapter, session.writer, "stream_read_error") diff --git a/backend/internal/service/antigravity_gateway_gemini.go b/backend/internal/service/antigravity_gateway_gemini.go index eb46a2d4d4e8..72480e889bb3 100644 --- a/backend/internal/service/antigravity_gateway_gemini.go +++ b/backend/internal/service/antigravity_gateway_gemini.go @@ -96,6 +96,9 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) return nil, s.writeGoogleError(c, http.StatusForbidden, fmt.Sprintf("model %s not in whitelist", originalModel)) } + if geminiImageGenerationDenied(c, body, originalModel, mappedModel) { + return nil, s.writeGoogleError(c, http.StatusBadRequest, ImageGenerationPermissionMessage()) + } forwardedModel := mappedModel // 获取 access_token diff --git a/backend/internal/service/antigravity_gateway_retry.go b/backend/internal/service/antigravity_gateway_retry.go index ea444fc755b5..1ce6b9b32420 100644 --- a/backend/internal/service/antigravity_gateway_retry.go +++ b/backend/internal/service/antigravity_gateway_retry.go @@ -858,10 +858,12 @@ func sleepAntigravityBackoffWithContext(ctx context.Context, attempt int) bool { } } -// isSingleAccountRetry 检查 context 中是否设置了单账号退避重试标记 +// isSingleAccountRetry 检查是否处于单账号退避重试模式：显式标记优先，否则使用请求挂载的延迟判定。 func isSingleAccountRetry(ctx context.Context) bool { - v, _ := SingleAccountRetryFromContext(ctx) - return v + if v, ok := SingleAccountRetryFromContext(ctx); ok { + return v + } + return singleAccountRetryProbeResult(ctx) } // setModelRateLimitByModelName 使用官方模型 ID 设置模型级限流 diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index b1611ae976b2..25478a57a9a1 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -884,6 +884,20 @@ func (s *BillingService) initFallbackPricing() { LongContextOutputMultiplier: 2, } + // xAI Grok 4.7 Fast (Cursor and Grok Build only): the same model on faster + // infrastructure at $4 input / $1 cached input / $12 output below 200k; + // long-context rates are $6 / $1.50 / $18 (1.5x of the Fast base card). + s.fallbackPrices["grok-4.7-fast"] = &ModelPricing{ + InputPricePerToken: 4e-6, + OutputPricePerToken: 12e-6, + CacheReadPricePerToken: 1e-6, + SupportsCacheBreakdown: false, + LongContextInputThreshold: 200000, + LongContextThresholdInclusive: true, + LongContextInputMultiplier: 1.5, + LongContextOutputMultiplier: 1.5, + } + // xAI Grok 4.3: $1.25 input / $0.20 cached / $2.50 output below 200k; // long-context rates are $2.50 / $0.40 / $5. s.fallbackPrices["grok-4.3"] = &ModelPricing{ @@ -1172,6 +1186,11 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing { return s.fallbackPrices["grok-4.6"] case "grok-4.7", "grok-4.7-latest": return s.fallbackPrices["grok-4.7"] + // Grok 4.7 Fast IDs (grok-4.7-fast, and the Grok Build ID + // grok-4.7-build-fast). Without these cases both fall into the unknown-Grok + // family card and bill at the standard 4.7 rate (#8119). + case "grok-4.7-fast", "grok-4.7-build-fast": + return s.fallbackPrices["grok-4.7-fast"] case "grok-4.5", "grok-4.5-latest": return s.fallbackPrices["grok-4.5"] case "grok-3-mini": diff --git a/backend/internal/service/billing_service_test.go b/backend/internal/service/billing_service_test.go index 45544508ec1c..d02c942cb3fc 100644 --- a/backend/internal/service/billing_service_test.go +++ b/backend/internal/service/billing_service_test.go @@ -1203,6 +1203,58 @@ func TestGetModelPricing_Grok47OfficialFallback(t *testing.T) { } } +func TestGetModelPricing_Grok47FastOfficialFallback(t *testing.T) { + svc := newTestBillingService() + + for _, model := range []string{"grok-4.7-fast", "grok-4.7-build-fast", "GROK-4.7-FAST"} { + model := model + t.Run(model, func(t *testing.T) { + pricing, err := svc.GetModelPricing(model) + require.NoError(t, err) + require.InDelta(t, 4e-6, pricing.InputPricePerToken, 1e-12) + require.InDelta(t, 12e-6, pricing.OutputPricePerToken, 1e-12) + require.InDelta(t, 1e-6, pricing.CacheReadPricePerToken, 1e-12) + require.Equal(t, 200000, pricing.LongContextInputThreshold) + require.True(t, pricing.LongContextThresholdInclusive) + require.InDelta(t, 1.5, pricing.LongContextInputMultiplier, 1e-12) + require.InDelta(t, 1.5, pricing.LongContextOutputMultiplier, 1e-12) + require.False(t, pricing.SupportsCacheBreakdown) + }) + } + require.True(t, svc.HasIdentifiedTokenPricing("grok-4.7-fast")) +} + +// Issue #8119: Grok 4.7 Fast used to inherit the unknown-Grok card and bill at +// the standard grok-4.7 rate. +func TestCalculateCost_Grok47FastUsesFastCard(t *testing.T) { + svc := newTestBillingService() + + short := UsageTokens{InputTokens: 100_000, CacheReadTokens: 50_000, OutputTokens: 10_000} + standard, err := svc.CalculateCost("grok-4.7", short, 1) + require.NoError(t, err) + fast, err := svc.CalculateCost("grok-4.7-fast", short, 1) + require.NoError(t, err) + require.InDelta(t, 0.4, fast.InputCost, 1e-12) + require.InDelta(t, 0.05, fast.CacheReadCost, 1e-12) + require.InDelta(t, 0.12, fast.OutputCost, 1e-12) + require.False(t, fast.LongContextBillingApplied) + // Below 200k the Fast card is exactly twice the standard card. + require.InDelta(t, standard.TotalCost*2, fast.TotalCost, 1e-12) + + // Reaching 200k prompt tokens moves the whole request to $6 / $1.50 / $18. + long := UsageTokens{InputTokens: 150_000, CacheReadTokens: 50_000, OutputTokens: 10_000} + standardLong, err := svc.CalculateCost("grok-4.7", long, 1) + require.NoError(t, err) + fastLong, err := svc.CalculateCost("grok-4.7-fast", long, 1) + require.NoError(t, err) + require.True(t, fastLong.LongContextBillingApplied) + require.InDelta(t, 0.9, fastLong.InputCost, 1e-12) + require.InDelta(t, 0.075, fastLong.CacheReadCost, 1e-12) + require.InDelta(t, 0.18, fastLong.OutputCost, 1e-12) + // Long-context Fast is 1.5x the standard long-context card. + require.InDelta(t, standardLong.TotalCost*1.5, fastLong.TotalCost, 1e-12) +} + func TestGetModelPricing_GrokOfficialFamilyCards(t *testing.T) { svc := newTestBillingService() for _, tc := range []struct { diff --git a/backend/internal/service/codex_fingerprint_account_view_test.go b/backend/internal/service/codex_fingerprint_account_view_test.go index 883da32631b1..0513f0aae7c5 100644 --- a/backend/internal/service/codex_fingerprint_account_view_test.go +++ b/backend/internal/service/codex_fingerprint_account_view_test.go @@ -87,7 +87,7 @@ func TestCodexFingerprintAccountViewCustomUAAndUnknownDeviceFields(t *testing.T) require.Contains(t, view.Identity.UserAgent, "0.162.0") account.Credentials["user_agent"] = "codex-tui/0.150.0" view = DescribeCodexFingerprintAccount(account, account) - require.Equal(t, "codex-tui/0.162.0", view.Identity.UserAgent) + require.Equal(t, "codex_cli_rs/0.162.0", view.Identity.UserAgent) require.NotNil(t, view.EffectiveDevice) require.Nil(t, view.EffectiveDevice.OSType) require.Nil(t, view.EffectiveDevice.Terminal) diff --git a/backend/internal/service/gateway_claude_prelude_test.go b/backend/internal/service/gateway_claude_prelude_test.go deleted file mode 100644 index f9592a82c21f..000000000000 --- a/backend/internal/service/gateway_claude_prelude_test.go +++ /dev/null @@ -1,206 +0,0 @@ -package service - -import ( - "context" - "errors" - "fmt" - "io" - "net/http" - "net/http/httptest" - "strings" - "testing" - "time" - - "github.com/gin-gonic/gin" - "github.com/stretchr/testify/require" -) - -const claudePreludeSignatureError = `{"type":"error","error":{"type":"invalid_request_error","message":"Invalid signature in thinking block"}}` - -func claudePreludeFrame(event, data string) string { - return "event: " + event + "\ndata: " + data + "\n\n" -} - -func claudePreludeStart(id string, inputTokens int) string { - return claudePreludeFrame("message_start", fmt.Sprintf(`{"type":"message_start","message":{"id":%q,"model":"claude-sonnet-5-5","usage":{"input_tokens":%d}}}`, id, inputTokens)) -} - -func claudePreludeTestContext() (*httptest.ResponseRecorder, *gin.Context) { - gin.SetMode(gin.TestMode) - recorder := httptest.NewRecorder() - c, _ := gin.CreateTestContext(recorder) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) - return recorder, c -} - -type claudePreludeFlushObserver struct { - gin.ResponseWriter - onFlush func() -} - -func (w *claudePreludeFlushObserver) Flush() { - w.ResponseWriter.Flush() - w.onFlush() -} - -func runClaudePreludeTestStream(t *testing.T, svc *GatewayService, c *gin.Context, stream string, recoverPrelude bool) (*streamingResult, error) { - t.Helper() - resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(stream))} - defer func() { _ = resp.Body.Close() }() - return svc.handleStreamingResponse(context.Background(), resp, c, &Account{ID: 1, Platform: PlatformAnthropic}, time.Now(), "claude-sonnet-5-5", "claude-sonnet-5-5", false, recoverPrelude) -} - -func TestClaudePrelude_SignatureRejectionDiscardsStartButKeepsAttemptUsage(t *testing.T) { - svc := newStreamingResponseTestGatewayService() - recorder, c := claudePreludeTestContext() - stream := claudePreludeStart("failed-message", 123) + claudePreludeFrame("ping", `{"type":"ping"}`) + ": keepalive\n\n" + claudePreludeFrame("error", claudePreludeSignatureError) - - result, err := runClaudePreludeTestStream(t, svc, c, stream, true) - var sseErr *sseStreamErrorEventError - require.True(t, errors.As(err, &sseErr)) - require.False(t, sseErr.MessageStarted, "ping and comments do not start this message") - require.NotNil(t, result) - require.Equal(t, 123, result.usage.InputTokens, "discarding a client prelude must not discard this attempt's observed metering") - require.NotContains(t, recorder.Body.String(), "failed-message") - require.Contains(t, recorder.Body.String(), "event: ping") - require.Contains(t, recorder.Body.String(), ": keepalive") - - // A successful repaired attempt contributes its own usage and one message ID. - beginUpstreamResponseModelObservation(c) - success := claudePreludeStart("repaired-message", 7) + claudePreludeFrame("message_delta", `{"type":"message_delta","usage":{"output_tokens":3}}`) + claudePreludeFrame("message_stop", `{"type":"message_stop"}`) - result, err = runClaudePreludeTestStream(t, svc, c, success, true) - require.NoError(t, err) - require.Equal(t, 7, result.usage.InputTokens) - require.Equal(t, 3, result.usage.OutputTokens) - require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: message_start")) - require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: message_stop")) - require.Contains(t, recorder.Body.String(), "repaired-message") -} - -func TestClaudePrelude_PingIsFlushedBeforeStartCommit(t *testing.T) { - svc := newStreamingResponseTestGatewayService() - recorder, c := claudePreludeTestContext() - beforeCommit := "" - c.Writer = &claudePreludeFlushObserver{ResponseWriter: c.Writer, onFlush: func() { - if beforeCommit == "" { - beforeCommit = recorder.Body.String() - } - }} - stream := claudePreludeStart("message", 5) + claudePreludeFrame("ping", `{"type":"ping"}`) + claudePreludeFrame("message_stop", `{"type":"message_stop"}`) - _, err := runClaudePreludeTestStream(t, svc, c, stream, true) - require.NoError(t, err) - require.Contains(t, beforeCommit, "event: ping") - require.NotContains(t, beforeCommit, "event: message_start") -} - -func TestClaudePrelude_CommitsStartBeforeContentAndPreservesSignatureDelta(t *testing.T) { - svc := newStreamingResponseTestGatewayService() - recorder, c := claudePreludeTestContext() - start := claudePreludeStart("message", 5) - contentStart := claudePreludeFrame("content_block_start", `{"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}`) - thinkingDelta := claudePreludeFrame("content_block_delta", `{"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"A thought"}}`) - signatureDelta := claudePreludeFrame("content_block_delta", `{"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"opaque+/=fragment"}}`) - contentStop := claudePreludeFrame("content_block_stop", `{"type":"content_block_stop","index":0}`) - stop := claudePreludeFrame("message_stop", `{"type":"message_stop"}`) - - _, err := runClaudePreludeTestStream(t, svc, c, start+contentStart+thinkingDelta+signatureDelta+contentStop+stop, true) - require.NoError(t, err) - require.Equal(t, start+contentStart+thinkingDelta+signatureDelta+contentStop+stop, recorder.Body.String(), "buffering must preserve all event order and signature bytes") -} - -func TestClaudePrelude_ContentOutputPreventsSignatureRecovery(t *testing.T) { - svc := newStreamingResponseTestGatewayService() - recorder, c := claudePreludeTestContext() - stream := claudePreludeStart("message", 5) + claudePreludeFrame("content_block_start", `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`) + claudePreludeFrame("error", claudePreludeSignatureError) - result, err := runClaudePreludeTestStream(t, svc, c, stream, true) - var sseErr *sseStreamErrorEventError - require.True(t, errors.As(err, &sseErr)) - require.True(t, sseErr.MessageStarted) - require.Equal(t, 5, result.usage.InputTokens) - require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: message_start")) - require.Contains(t, recorder.Body.String(), "event: content_block_start") -} - -func TestClaudePrelude_OversizedStartCommitsAndPreventsRecovery(t *testing.T) { - svc := newStreamingResponseTestGatewayService() - recorder, c := claudePreludeTestContext() - start := claudePreludeFrame("message_start", fmt.Sprintf(`{"type":"message_start","message":{"id":%q,"usage":{"input_tokens":5}}}`, strings.Repeat("x", claudeMessageStartPreludeLimit))) - _, err := runClaudePreludeTestStream(t, svc, c, start+claudePreludeFrame("error", claudePreludeSignatureError), true) - var sseErr *sseStreamErrorEventError - require.True(t, errors.As(err, &sseErr)) - require.True(t, sseErr.MessageStarted) - require.Equal(t, start, recorder.Body.String()) -} - -func TestClaudePrelude_NonSignatureRejectionCommitsStart(t *testing.T) { - for _, errorJSON := range []string{ - `{"type":"error","error":{"type":"service_error","message":"Service temporarily unavailable"}}`, - `{"type":"error","error":{"type":"service_error","message":"Invalid signature in thinking block"}}`, - `{"type":"error","error":{"type":"permission_error","message":"Invalid signature in thinking block"}}`, - } { - t.Run(errorJSON, func(t *testing.T) { - svc := newStreamingResponseTestGatewayService() - recorder, c := claudePreludeTestContext() - start := claudePreludeStart("message", 5) - result, err := runClaudePreludeTestStream(t, svc, c, start+claudePreludeFrame("error", errorJSON), true) - var sseErr *sseStreamErrorEventError - require.True(t, errors.As(err, &sseErr)) - require.True(t, sseErr.MessageStarted) - require.Equal(t, 5, result.usage.InputTokens) - require.Equal(t, start, recorder.Body.String()) - }) - } -} - -func TestClaudePrelude_EOFCommitsStartAndEndsWithSingleError(t *testing.T) { - svc := newStreamingResponseTestGatewayService() - recorder, c := claudePreludeTestContext() - start := claudePreludeStart("message", 5) - result, err := runClaudePreludeTestStream(t, svc, c, start, true) - require.ErrorContains(t, err, "missing terminal event") - require.Equal(t, 5, result.usage.InputTokens) - require.True(t, strings.HasPrefix(recorder.Body.String(), start)) - require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: error")) - require.True(t, IsResponseCommitted(c), "the handler must not append a second terminal after the prelude EOF error") -} - -func TestClaudePrelude_DisabledRecoveryDoesNotHoldStart(t *testing.T) { - svc := newStreamingResponseTestGatewayService() - recorder, c := claudePreludeTestContext() - start := claudePreludeStart("message", 5) - _, err := runClaudePreludeTestStream(t, svc, c, start+claudePreludeFrame("error", claudePreludeSignatureError), false) - var sseErr *sseStreamErrorEventError - require.True(t, errors.As(err, &sseErr)) - require.True(t, sseErr.MessageStarted) - require.Equal(t, start, recorder.Body.String()) -} - -func TestClaudePrelude_RepairTimeoutSkipsHealthPenaltyAndNormalTimeoutKeepsIt(t *testing.T) { - for _, suppressPenalty := range []bool{false, true} { - t.Run(fmt.Sprintf("suppress_%t", suppressPenalty), func(t *testing.T) { - svc := newStreamingResponseTestGatewayService() - svc.cfg.Gateway.StreamDataIntervalTimeout = 1 - repo := &gatewayForwardErrorPolicyRepoStub{} - settings := NewSettingService(&gatewayTTLSettingRepo{data: map[string]string{ - SettingKeyStreamTimeoutSettings: `{"enabled":true,"action":"temp_unsched","threshold_count":1,"threshold_window_minutes":1,"temp_unsched_minutes":1}`, - }}, svc.cfg) - svc.rateLimitService = NewRateLimitService(repo, nil, svc.cfg, nil, nil) - svc.rateLimitService.SetSettingService(settings) - recorder, c := claudePreludeTestContext() - pr, pw := io.Pipe() - defer func() { _ = pr.Close() }() - defer func() { _ = pw.Close() }() - resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: pr} - - _, err := svc.handleStreamingResponse(context.Background(), resp, c, &Account{ID: 1}, time.Now(), "model", "model", false, false, suppressPenalty) - require.ErrorContains(t, err, "stream data interval timeout") - if suppressPenalty { - require.Zero(t, repo.tempCalls, "a failed signature repair must not park the account") - } else { - require.Equal(t, 1, repo.tempCalls, "normal streaming timeout keeps its existing health action") - } - require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: error")) - require.Contains(t, recorder.Body.String(), `"type":"stream_timeout"`) - }) - } -} diff --git a/backend/internal/service/gateway_claude_signature_recovery.go b/backend/internal/service/gateway_claude_signature_recovery.go index d535904899ae..0aff6e99a7e4 100644 --- a/backend/internal/service/gateway_claude_signature_recovery.go +++ b/backend/internal/service/gateway_claude_signature_recovery.go @@ -1,7 +1,6 @@ package service import ( - "bytes" "context" "crypto/sha256" "encoding/hex" @@ -120,7 +119,7 @@ func (s *GatewayService) claudeSignatureErrorCandidate(ctx context.Context, acco return s.settingService != nil && s.isSignatureErrorPattern(ctx, account, body) } -// run provides one first-stage repair shared by HTTP and SSE rejection paths. +// run provides the first-stage repair for HTTP signature rejections. // Its decision budget begins at the rejection, not before the first upstream // request. Normal transport/stream timeouts remain in force after sending. func (s *claudeSignatureRecoveryState) run( @@ -133,22 +132,11 @@ func (s *claudeSignatureRecoveryState) run( diagnostic.Reason = "context_canceled" return nil, nil, nil } - if diagnostic.MessageStarted { - diagnostic.Reason = "output_started" - return nil, nil, nil - } if s.used { diagnostic.Reason = "attempt_limit" return nil, nil, nil } filteredBody := FilterThinkingBlocksForRetry(baseBody, model) - // The existing HTTP two-stage fallback can need a no-op first stage - // before a tool-specific rejection authorizes its second-stage filter. - // New SSE recovery must not replay an unchanged body. - if bytes.Equal(filteredBody, baseBody) && diagnostic.Source == "sse_error" { - diagnostic.Reason = "body_unchanged" - return nil, nil, nil - } s.used = true req, wireBody, err := build(filteredBody) if err != nil { diff --git a/backend/internal/service/gateway_claude_signature_recovery_test.go b/backend/internal/service/gateway_claude_signature_recovery_test.go index e9d77ad0e624..4f4dda188efc 100644 --- a/backend/internal/service/gateway_claude_signature_recovery_test.go +++ b/backend/internal/service/gateway_claude_signature_recovery_test.go @@ -95,106 +95,53 @@ func claudeRecoveryService(upstream *claudeRecoverySequenceUpstream, enabled boo settingService: NewSettingService(&gatewayTTLSettingRepo{data: map[string]string{SettingKeyRectifierSettings: settings}}, cfg)} } -func TestClaudeSignatureRecoveryForwardHTTPAndSSE(t *testing.T) { - for _, tc := range []struct { - name string - status int - prelude string - }{ - {"http_400", 400, ""}, - {"sse_before_message", 200, ""}, - {"sse_after_ping", 200, "event: ping\ndata: {\"type\":\"ping\"}\n\n"}, - {"sse_after_held_message_start", 200, "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"rejected-message\",\"model\":\"other-model\",\"usage\":{\"input_tokens\":900}}}\n\n"}, - } { - t.Run(tc.name, func(t *testing.T) { - gin.SetMode(gin.TestMode) - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) - firstBody := claudeRecoverySignatureError - if tc.status == 200 { - firstBody = tc.prelude + "event: error\ndata: " + firstBody + "\n\n" - } - upstream := &claudeRecoverySequenceUpstream{responses: []*http.Response{ - claudeRecoveryResponse(tc.status, firstBody, "initial-rid"), claudeRecoveryResponse(200, claudeRecoverySuccessStream(), "repaired-rid"), - }} - initialBody := &claudeRecoveryCloseObservedBody{ReadCloser: upstream.responses[0].Body} - upstream.responses[0].Body = initialBody - upstream.beforeSend = func(index int) { - if index == 1 { - require.True(t, initialBody.closed, "the rejected response must release its connection before repair sends") - } - } - account := newAnthropicAPIKeyAccountForTest() - account.Extra = nil - parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(claudeRecoveryTestBody)), PlatformAnthropic) - require.NoError(t, err) - acceptedCalls := 0 - parsed.OnUpstreamAccepted = func() { acceptedCalls++ } - result, err := claudeRecoveryService(upstream, true).Forward(context.Background(), c, account, parsed) - require.NoError(t, err) - require.NotNil(t, result) - require.Len(t, upstream.requests, 2) - require.Equal(t, []int64{account.ID, account.ID}, upstream.accountIDs) - require.Equal(t, upstream.requests[0].URL, upstream.requests[1].URL) - require.Equal(t, upstream.requests[0].Header.Get("X-Api-Key"), upstream.requests[1].Header.Get("X-Api-Key")) - require.Equal(t, 1, acceptedCalls) - require.Equal(t, 7, result.Usage.InputTokens) - require.Equal(t, 3, result.Usage.OutputTokens) - require.Equal(t, "repaired-rid", result.RequestID) - require.False(t, result.UpstreamResponseModelConflict) - require.NotContains(t, rec.Body.String(), "rejected-message") - require.Equal(t, 1, strings.Count(rec.Body.String(), "event: message_start")) - require.NotContains(t, string(upstream.bodies[1]), "private-signature-value") - require.Equal(t, "between_tools", gjson.GetBytes(upstream.bodies[1], "thinking.type").String()) - require.Equal(t, "high", gjson.GetBytes(upstream.bodies[1], "output_config.effort").String()) - require.True(t, bytes.Equal(parsed.Body.Bytes(), upstream.bodies[1])) - trigger := claudeRecoveryOpsEvents(t, c)[0] - require.NotNil(t, trigger.SignatureRecovery) - require.Equal(t, "accepted", trigger.SignatureRecovery.Outcome) - require.Equal(t, 1, trigger.SignatureRecovery.Attempts) - require.Equal(t, 1, trigger.SignatureRecovery.Inbound.TotalCount) - diagnosticJSON, err := json.Marshal(trigger.SignatureRecovery) - require.NoError(t, err) - require.NotContains(t, string(diagnosticJSON), "private-signature-value") - require.NotContains(t, string(diagnosticJSON), "old reasoning") - }) - } -} - -func TestClaudeSignatureRecoveryDoesNotReplaySSEAgainOrAfterOutput(t *testing.T) { - for _, prefix := range []string{"", "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":4}}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"partial\"}}\n\n"} { - t.Run(fmt.Sprintf("output_%t", prefix != ""), func(t *testing.T) { - gin.SetMode(gin.TestMode) - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) - errStream := prefix + "event: error\ndata: " + claudeRecoverySignatureError + "\n\n" - upstream := &claudeRecoverySequenceUpstream{responses: []*http.Response{claudeRecoveryResponse(200, errStream, "one"), claudeRecoveryResponse(200, errStream, "two")}} - account := newAnthropicAPIKeyAccountForTest() - account.Extra = nil - parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(claudeRecoveryTestBody)), PlatformAnthropic) - require.NoError(t, err) - _, err = claudeRecoveryService(upstream, true).Forward(context.Background(), c, account, parsed) - require.Error(t, err) - var failure *UpstreamFailoverError - require.ErrorAs(t, err, &failure) - require.Equal(t, 400, failure.StatusCode) - require.False(t, failure.ShouldRetryNextAccount()) - require.False(t, failure.ShouldReportAccountScheduleFailure()) - wantRequests := 2 - if prefix != "" { - wantRequests = 1 - } - require.Len(t, upstream.requests, wantRequests) - last := claudeRecoveryOpsEvents(t, c)[wantRequests-1] - wantReason := "attempt_limit" - if prefix != "" { - wantReason = "output_started" - } - require.Equal(t, wantReason, last.SignatureRecovery.Reason) - }) +func TestClaudeSignatureRecoveryForwardHTTP(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + upstream := &claudeRecoverySequenceUpstream{responses: []*http.Response{ + claudeRecoveryResponse(400, claudeRecoverySignatureError, "initial-rid"), claudeRecoveryResponse(200, claudeRecoverySuccessStream(), "repaired-rid"), + }} + initialBody := &claudeRecoveryCloseObservedBody{ReadCloser: upstream.responses[0].Body} + upstream.responses[0].Body = initialBody + upstream.beforeSend = func(index int) { + if index == 1 { + require.True(t, initialBody.closed, "the rejected response must release its connection before repair sends") + } } + account := newAnthropicAPIKeyAccountForTest() + account.Extra = nil + parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(claudeRecoveryTestBody)), PlatformAnthropic) + require.NoError(t, err) + acceptedCalls := 0 + parsed.OnUpstreamAccepted = func() { acceptedCalls++ } + result, err := claudeRecoveryService(upstream, true).Forward(context.Background(), c, account, parsed) + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.requests, 2) + require.Equal(t, []int64{account.ID, account.ID}, upstream.accountIDs) + require.Equal(t, upstream.requests[0].URL, upstream.requests[1].URL) + require.Equal(t, upstream.requests[0].Header.Get("X-Api-Key"), upstream.requests[1].Header.Get("X-Api-Key")) + require.Equal(t, 1, acceptedCalls) + require.Equal(t, 7, result.Usage.InputTokens) + require.Equal(t, 3, result.Usage.OutputTokens) + require.Equal(t, "repaired-rid", result.RequestID) + require.False(t, result.UpstreamResponseModelConflict) + require.Equal(t, 1, strings.Count(rec.Body.String(), "event: message_start")) + require.NotContains(t, string(upstream.bodies[1]), "private-signature-value") + require.Equal(t, "between_tools", gjson.GetBytes(upstream.bodies[1], "thinking.type").String()) + require.Equal(t, "high", gjson.GetBytes(upstream.bodies[1], "output_config.effort").String()) + require.True(t, bytes.Equal(parsed.Body.Bytes(), upstream.bodies[1])) + trigger := claudeRecoveryOpsEvents(t, c)[0] + require.NotNil(t, trigger.SignatureRecovery) + require.Equal(t, "accepted", trigger.SignatureRecovery.Outcome) + require.Equal(t, 1, trigger.SignatureRecovery.Attempts) + require.Equal(t, 1, trigger.SignatureRecovery.Inbound.TotalCount) + diagnosticJSON, err := json.Marshal(trigger.SignatureRecovery) + require.NoError(t, err) + require.NotContains(t, string(diagnosticJSON), "private-signature-value") + require.NotContains(t, string(diagnosticJSON), "old reasoning") } func TestClaudeSignatureRecoveryDisabledAndPassthrough(t *testing.T) { @@ -220,7 +167,7 @@ func TestClaudeSignatureRecoveryDisabledAndPassthrough(t *testing.T) { } func TestClaudeSignatureRecoveryDecisionAndFailureEvidence(t *testing.T) { - for _, reason := range []string{"build_failed", "transport_failed", "empty_response", "budget_exhausted", "context_canceled", "body_unchanged", "output_started", "attempt_limit"} { + for _, reason := range []string{"build_failed", "transport_failed", "empty_response", "budget_exhausted", "context_canceled", "attempt_limit"} { t.Run(reason, func(t *testing.T) { state := newClaudeSignatureRecoveryState([]byte(claudeRecoveryTestBody)) diag := state.diagnostic("http_400", false) @@ -234,13 +181,6 @@ func TestClaudeSignatureRecoveryDecisionAndFailureEvidence(t *testing.T) { ctx, cancel = context.WithCancel(ctx) cancel() } - if reason == "body_unchanged" { - diag.Source = "sse_error" - body = []byte(`{"messages":[{"role":"user","content":"hi"}]}`) - } - if reason == "output_started" { - diag.MessageStarted = true - } if reason == "attempt_limit" { state.used = true } @@ -303,82 +243,6 @@ func TestClaudeSignatureRecoveryFingerprintAndSanitizer(t *testing.T) { require.NotContains(t, string(encoded), "private text") } -func TestClaudeSignatureRecoverySharedHTTPAndSSELimit(t *testing.T) { - gin.SetMode(gin.TestMode) - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) - errStream := "event: error\ndata: " + claudeRecoverySignatureError + "\n\n" - upstream := &claudeRecoverySequenceUpstream{responses: []*http.Response{ - claudeRecoveryResponse(400, claudeRecoverySignatureError, "http-rejection"), - claudeRecoveryResponse(200, errStream, "sse-rejection"), - }} - account := newAnthropicAPIKeyAccountForTest() - account.Extra = nil - parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(claudeRecoveryTestBody)), PlatformAnthropic) - require.NoError(t, err) - _, err = claudeRecoveryService(upstream, true).Forward(context.Background(), c, account, parsed) - require.Error(t, err) - require.Len(t, upstream.requests, 2) - attempts := claudeRecoveryOpsEvents(t, c) - require.Equal(t, "attempt_limit", attempts[len(attempts)-1].SignatureRecovery.Reason) -} - -func TestClaudeSignatureRecoveryFailureStopsWithoutHealthPenalty(t *testing.T) { - for _, streamFailure := range []bool{false, true} { - t.Run(fmt.Sprintf("stream_failure_%t", streamFailure), func(t *testing.T) { - gin.SetMode(gin.TestMode) - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) - serviceError := `{"type":"error","error":{"type":"service_error","message":"temporarily unavailable"}}` - repairedStatus := 503 - repairedBody := serviceError - if streamFailure { - repairedStatus = 200 - repairedBody = "event: error\ndata: " + serviceError + "\n\n" - } - upstream := &claudeRecoverySequenceUpstream{responses: []*http.Response{ - claudeRecoveryResponse(200, "event: error\ndata: "+claudeRecoverySignatureError+"\n\n", "initial"), - claudeRecoveryResponse(repairedStatus, repairedBody, "retry"), - }} - account := newAnthropicAPIKeyAccountForTest() - account.Extra = nil - parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(claudeRecoveryTestBody)), PlatformAnthropic) - require.NoError(t, err) - _, err = claudeRecoveryService(upstream, true).Forward(context.Background(), c, account, parsed) - require.Error(t, err) - var failure *UpstreamFailoverError - require.ErrorAs(t, err, &failure) - require.False(t, failure.ShouldRetryNextAccount()) - require.False(t, failure.ShouldReportAccountScheduleFailure()) - require.Equal(t, serviceError, string(failure.ResponseBody)) - require.Len(t, upstream.requests, 2) - }) - } -} - -func TestClaudeSignatureRecoveryExplicitRequestPatternMatchesPrelude(t *testing.T) { - gin.SetMode(gin.TestMode) - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) - // This is an existing HTTP rectifier match without a "thinking" word; it - // must agree with the prelude handler when the provider declares HTTP 400. - first := "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"discard-me\"}}\n\nevent: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"signature is invalid\"}}\n\n" - upstream := &claudeRecoverySequenceUpstream{responses: []*http.Response{claudeRecoveryResponse(200, first, "first"), claudeRecoveryResponse(200, claudeRecoverySuccessStream(), "second")}} - account := newAnthropicAPIKeyAccountForTest() - account.Extra = nil - parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(claudeRecoveryTestBody)), PlatformAnthropic) - require.NoError(t, err) - result, err := claudeRecoveryService(upstream, true).Forward(context.Background(), c, account, parsed) - require.NoError(t, err) - require.NotNil(t, result) - require.Len(t, upstream.requests, 2) - require.NotContains(t, rec.Body.String(), "discard-me") - require.Equal(t, 1, strings.Count(rec.Body.String(), "event: message_start")) -} - func TestClaudeSignatureRecoverySlowInitialResponseDoesNotSpendDecisionBudget(t *testing.T) { clock := time.Unix(0, 0) state := newClaudeSignatureRecoveryState([]byte(claudeRecoveryTestBody)) @@ -405,41 +269,6 @@ func TestClaudeSignatureRecoverySlowInitialResponseDoesNotSpendDecisionBudget(t require.Equal(t, int64(1000), diag.ElapsedMs) } -func TestClaudeSignatureRecoveryIncompleteRepairedStreamRecordsActualAttempt(t *testing.T) { - for _, initialHTTP := range []bool{false, true} { - t.Run(fmt.Sprintf("initial_http_%t", initialHTTP), func(t *testing.T) { - gin.SetMode(gin.TestMode) - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) - firstStatus := 200 - firstBody := "event: error\ndata: " + claudeRecoverySignatureError + "\n\n" - if initialHTTP { - firstStatus = 400 - firstBody = claudeRecoverySignatureError - } - incomplete := strings.Split(claudeRecoverySuccessStream(), "event: message_stop")[0] - upstream := &claudeRecoverySequenceUpstream{responses: []*http.Response{claudeRecoveryResponse(firstStatus, firstBody, "initial"), claudeRecoveryResponse(200, incomplete, "incomplete-retry")}} - account := newAnthropicAPIKeyAccountForTest() - account.Extra = nil - parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(claudeRecoveryTestBody)), PlatformAnthropic) - require.NoError(t, err) - _, err = claudeRecoveryService(upstream, true).Forward(context.Background(), c, account, parsed) - require.Error(t, err) - require.Contains(t, err.Error(), "missing terminal") - attempts := claudeRecoveryOpsEvents(t, c) - require.Equal(t, "response_read_failed", attempts[0].SignatureRecovery.Outcome) - require.Equal(t, "stream_incomplete", attempts[0].SignatureRecovery.Reason) - last := attempts[len(attempts)-1] - require.Equal(t, "stream_failure", last.Kind) - require.Equal(t, 502, last.UpstreamStatusCode) - require.Equal(t, 200, last.UpstreamHTTPStatusCode) - require.Equal(t, "incomplete-retry", last.UpstreamRequestID) - require.Len(t, upstream.requests, 2) - }) - } -} - func TestClaudeSignatureRecoveryHTTPToolOnlyKeepsExistingSecondStage(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() diff --git a/backend/internal/service/gateway_claude_stream_error.go b/backend/internal/service/gateway_claude_stream_error.go deleted file mode 100644 index cf548f968de0..000000000000 --- a/backend/internal/service/gateway_claude_stream_error.go +++ /dev/null @@ -1,122 +0,0 @@ -package service - -import ( - "net/http" - "strings" - - "github.com/tidwall/gjson" -) - -// ClaudeStreamErrorClassification describes an error carried inside an -// otherwise successful upstream SSE response. StatusCode is its semantic error -// status; it is not the HTTP status of the response carrying the event. -type ClaudeStreamErrorClassification struct { - StatusCode int - ClientStatusCode int - ClientErrorType string - ClientErrorCode string - ClientMessage string - RequestError bool -} - -// ClassifyClaudeStreamError classifies the original event without using whether -// client output has started. Output safety and signature recovery are separate -// decisions made by the caller. -func ClassifyClaudeStreamError(body []byte) ClaudeStreamErrorClassification { - errType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.type").String())) - result := ClaudeStreamErrorClassification{ - StatusCode: http.StatusBadGateway, - ClientStatusCode: http.StatusBadGateway, - ClientErrorType: "upstream_error", - ClientMessage: sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(body))), - } - if code := gjson.GetBytes(body, "error.code"); code.Type == gjson.String { - result.ClientErrorCode = strings.TrimSpace(code.String()) - } - switch errType { - case "invalid_request_error": - result.StatusCode = http.StatusBadRequest - result.RequestError = true - case "authentication_error": - result.StatusCode = http.StatusUnauthorized - case "billing_error": - result.StatusCode = http.StatusPaymentRequired - case "permission_error": - result.StatusCode = http.StatusForbidden - case "not_found_error": - result.StatusCode = http.StatusNotFound - result.RequestError = true - case "conflict_error": - result.StatusCode = http.StatusConflict - result.RequestError = true - case "request_too_large", "request_too_large_error": - result.StatusCode = http.StatusRequestEntityTooLarge - result.RequestError = true - case "rate_limit_error": - result.StatusCode = http.StatusTooManyRequests - case "api_error", "service_error": - result.StatusCode = http.StatusInternalServerError - case "overloaded_error": - result.StatusCode = 529 - case "timeout_error": - result.StatusCode = http.StatusGatewayTimeout - default: - errType = "upstream_error" - if IsClaudeThinkingSignatureError(body) { - result.StatusCode = http.StatusBadRequest - result.RequestError = true - errType = "invalid_request_error" - } - } - result.ClientErrorType = errType - if result.RequestError || result.StatusCode == http.StatusTooManyRequests { - result.ClientStatusCode = result.StatusCode - } else if result.StatusCode == 529 { - result.ClientStatusCode = http.StatusServiceUnavailable - } - if result.ClientMessage == "" { - result.ClientMessage = "Upstream returned an error in the response stream" - } - return result -} - -// FailoverError retains the original event for rules and administrator -// diagnostics. Request validation failures must not be replayed on another -// account or counted as account health failures. -func (classification ClaudeStreamErrorClassification) FailoverError(body []byte) *UpstreamFailoverError { - err := &UpstreamFailoverError{ - StatusCode: classification.StatusCode, - ResponseBody: body, - ClientStatusCode: classification.ClientStatusCode, - ClientErrorType: classification.ClientErrorType, - ClientErrorCode: classification.ClientErrorCode, - ClientMessage: classification.ClientMessage, - } - if classification.RequestError { - err.NextAccountAction = NextAccountStop - err.SuppressAccountHealthPenalty = true - } - return err -} - -// IsClaudeThinkingSignatureError recognizes explicit rejection of a thinking -// signature. Merely mentioning a signature or thinking is not sufficient: such -// messages can describe unrelated provider failures or user content. -func IsClaudeThinkingSignatureError(body []byte) bool { - message := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(body))) - if !strings.Contains(message, "signature") || - (!strings.Contains(message, "thinking") && !strings.Contains(message, "thought_signature")) { - return false - } - for _, rejection := range []string{ - "invalid", "not valid", "missing", "required", "malformed", "mismatch", - "does not match", "doesn't match", "cannot verify", "cannot be verified", - "unable to verify", "verification failed", "failed to verify", "failed to validate", - "different conversation", "another conversation", "different session", "another session", - } { - if strings.Contains(message, rejection) { - return true - } - } - return false -} diff --git a/backend/internal/service/gateway_claude_stream_error_test.go b/backend/internal/service/gateway_claude_stream_error_test.go deleted file mode 100644 index 5864814f161f..000000000000 --- a/backend/internal/service/gateway_claude_stream_error_test.go +++ /dev/null @@ -1,89 +0,0 @@ -package service - -import ( - "encoding/json" - "testing" - - "github.com/stretchr/testify/require" -) - -func TestClaudeClassifyStreamError(t *testing.T) { - tests := []struct { - name string - errorType string - message string - status int - clientStatus int - clientType string - requestError bool - }{ - {"invalid request", "invalid_request_error", "Invalid request body", 400, 400, "invalid_request_error", true}, - {"authentication", "authentication_error", "Authentication failed", 401, 502, "authentication_error", false}, - {"billing", "billing_error", "Payment required", 402, 502, "billing_error", false}, - {"permission", "permission_error", "Permission denied", 403, 502, "permission_error", false}, - {"not found", "not_found_error", "Model not found", 404, 404, "not_found_error", true}, - {"conflict", "conflict_error", "Request conflict", 409, 409, "conflict_error", true}, - {"too large", "request_too_large", "Request too large", 413, 413, "request_too_large", true}, - {"too large alias", "request_too_large_error", "Request too large", 413, 413, "request_too_large_error", true}, - {"rate limited", "rate_limit_error", "Rate limit exceeded", 429, 429, "rate_limit_error", false}, - {"api error", "api_error", "Internal error", 500, 502, "api_error", false}, - {"third party service", "service_error", "Service temporarily unavailable", 500, 502, "service_error", false}, - {"timeout", "timeout_error", "Request timed out", 504, 502, "timeout_error", false}, - {"overloaded", "overloaded_error", "Overloaded", 529, 503, "overloaded_error", false}, - {"unknown", "provider_specific_error", "Provider failure", 502, 502, "upstream_error", false}, - {"missing type signature", "", "Invalid `signature` in `thinking` block", 400, 400, "invalid_request_error", true}, - {"nil type signature", "", "Invalid `signature` in `thinking` block", 400, 400, "invalid_request_error", true}, - {"unknown type signature", "provider_specific_error", "Invalid `signature` in `thinking` block", 400, 400, "invalid_request_error", true}, - {"explicit permission beats signature text", "permission_error", "Invalid `signature` in `thinking` block", 403, 502, "permission_error", false}, - {"explicit authentication beats signature text", "authentication_error", "Invalid `signature` in `thinking` block", 401, 502, "authentication_error", false}, - {"explicit service beats signature text", "service_error", "Invalid `signature` in `thinking` block", 500, 502, "service_error", false}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - body, err := json.Marshal(map[string]any{"type": "error", "error": map[string]string{"type": test.errorType, "message": test.message, "code": "provider_code"}}) - require.NoError(t, err) - classification := ClassifyClaudeStreamError(body) - require.Equal(t, test.status, classification.StatusCode) - require.Equal(t, test.clientStatus, classification.ClientStatusCode) - require.Equal(t, test.clientType, classification.ClientErrorType) - require.Equal(t, test.message, classification.ClientMessage) - require.Equal(t, "provider_code", classification.ClientErrorCode) - require.Equal(t, test.requestError, classification.RequestError) - failoverErr := classification.FailoverError(body) - require.Equal(t, !test.requestError, failoverErr.ShouldRetryNextAccount()) - require.Equal(t, test.requestError, failoverErr.SuppressAccountHealthPenalty) - require.False(t, failoverErr.RetryableOnSameAccount) - require.Equal(t, body, failoverErr.ResponseBody) - }) - } -} - -func TestClaudeThinkingSignatureError(t *testing.T) { - for _, test := range []struct { - message string - want bool - }{ - {"Invalid `signature` in `thinking` block", true}, - {"messages.2.content.0.thinking.signature: Field required", true}, - {"thinking signature is bound to a different conversation", true}, - {"thought_signature verification failed", true}, - {"Invalid API key signature", false}, - {"thinking signature service is temporarily unavailable", false}, - {"The response uses thinking and a signature", false}, - {"all messages must have non-empty content", false}, - } { - t.Run(test.message, func(t *testing.T) { - body, err := json.Marshal(map[string]any{"error": map[string]string{"message": test.message}}) - require.NoError(t, err) - require.Equal(t, test.want, IsClaudeThinkingSignatureError(body)) - }) - } -} - -func TestClaudeClassifyStreamError_MalformedBody(t *testing.T) { - classification := ClassifyClaudeStreamError([]byte("not JSON")) - require.Equal(t, 502, classification.StatusCode) - require.Equal(t, "upstream_error", classification.ClientErrorType) - require.NotEmpty(t, classification.ClientMessage) - require.False(t, classification.RequestError) -} diff --git a/backend/internal/service/gateway_forward.go b/backend/internal/service/gateway_forward.go index ea24b1ba0c7f..05888bd83202 100644 --- a/backend/internal/service/gateway_forward.go +++ b/backend/internal/service/gateway_forward.go @@ -14,6 +14,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/claude" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/tidwall/gjson" "github.com/gin-gonic/gin" ) @@ -402,10 +403,6 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } // 重试循环 - // Keep the unsigned request view for signature recovery, even after an - // accepted attempt updates ParsedRequest to its actual signed wire body. - signatureRecoveryBase := body - var activeRecoveryDiagnostic *ClaudeSignatureRecoveryDiagnostic var resp *http.Response lastWireBody := body retryStart := time.Now() @@ -495,7 +492,6 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A }) if retryResp != nil && retryErr == nil { if retryResp.StatusCode < 400 { - activeRecoveryDiagnostic = recoveryDiagnostic // 重试请求被上游接受后同步 ParsedRequest，保证 usage/日志看到真实请求体。 lastWireBody = retryWireBody if err := replaceBody(retryWireBody); err != nil { @@ -546,7 +542,6 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A if retryErr2 == nil { recoveryDiagnostic.Outcome = "upstream_rejected" if retryResp2.StatusCode < 400 { - activeRecoveryDiagnostic = recoveryDiagnostic recoveryDiagnostic.Outcome = "accepted" // 二阶段工具块降级成功时也必须更新当前 body。 lastWireBody = retryWireBody2 @@ -856,135 +851,23 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A var firstTokenMs *int var clientDisconnect bool if reqStream { - allowSignatureRecovery := s.claudeSignatureRecoveryEnabled(ctx, account, reqModel) - for { - streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, reqModel, shouldMimicClaudeCode, - allowSignatureRecovery && !signatureRecovery.used, activeRecoveryDiagnostic != nil) - if err != nil { - var sseErr *sseStreamErrorEventError - if errors.As(err, &sseErr) { - streamErrorBody := []byte(sseErr.RawData) - classification := ClassifyClaudeStreamError(streamErrorBody) - if activeRecoveryDiagnostic != nil { - activeRecoveryDiagnostic.Outcome = "upstream_rejected" - } - var recoveryDiagnostic *ClaudeSignatureRecoveryDiagnostic - if classification.StatusCode == http.StatusBadRequest && s.claudeSignatureErrorCandidate(ctx, account, streamErrorBody) { - recoveryDiagnostic = signatureRecovery.diagnostic("sse_error", sseErr.MessageStarted) - if !allowSignatureRecovery { - recoveryDiagnostic.Reason = "disabled" - } - } - upstreamDetail := "" - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes - if maxBytes <= 0 { - maxBytes = 2048 - } - upstreamDetail = truncateString(sseErr.RawData, maxBytes) - } - - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - ProxyID: opsUpstreamProxyID(account), - ProxyName: opsUpstreamProxyName(account), - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: classification.StatusCode, - UpstreamHTTPStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "stream_error", - Message: classification.ClientMessage, - Detail: upstreamDetail, - SignatureRecovery: recoveryDiagnostic, - }) - if recoveryDiagnostic != nil && allowSignatureRecovery { - // Release the rejected response before acquiring another - // connection, including pools with a single host slot. - _ = resp.Body.Close() - retryResp, retryWireBody, recoveryErr := signatureRecovery.run(ctx, recoveryDiagnostic, signatureRecoveryBase, reqModel, - func(filteredBody []byte) (*http.Request, []byte, error) { - retryCtx, releaseRetryCtx := detachStreamUpstreamContext(ctx, reqStream) - defer releaseRetryCtx() - return s.buildUpstreamRequest(retryCtx, c, account, filteredBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) - }, - func(req *http.Request) (*http.Response, error) { - return s.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, tlsProfile) - }) - if recoveryErr != nil { - logger.LegacyPrintf("service.gateway", "Account %d: SSE signature recovery failed (%s): %v", account.ID, recoveryDiagnostic.Reason, recoveryErr) - } - if retryResp != nil { - _ = resp.Body.Close() - resp = retryResp - if resp.StatusCode < 400 { - if err := replaceBody(retryWireBody); err != nil { - return nil, err - } - beginUpstreamResponseModelObservation(c) - activeRecoveryDiagnostic = recoveryDiagnostic - continue - } - // A repair can fail with a real HTTP response, including while - // earlier keepalive pings have committed HTTP 200. Let the - // streaming-aware handler write the terminal, never append JSON. - retryBody, readErr := s.readUpstreamErrorBody(resp) - if readErr != nil { - recoveryDiagnostic.Outcome, recoveryDiagnostic.Reason = "response_read_failed", "response_read_failed" - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - ProxyID: opsUpstreamProxyID(account), ProxyName: opsUpstreamProxyName(account), - Platform: account.Platform, AccountID: account.ID, AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "http_error", Message: extractUpstreamErrorMessage(retryBody), - }) - failure := &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: retryBody, NextAccountAction: NextAccountStop, SuppressAccountHealthPenalty: true} - if resp.StatusCode == http.StatusBadRequest || resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusRequestEntityTooLarge { - failure.ClientStatusCode = resp.StatusCode - failure.ClientErrorType = "invalid_request_error" - failure.ClientMessage = sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(retryBody)) - } - return nil, failure - } - } - if activeRecoveryDiagnostic == nil && !sseErr.MessageStarted && classification.StatusCode == 529 { - syntheticResp := &http.Response{StatusCode: classification.StatusCode, Header: resp.Header.Clone(), Body: io.NopCloser(bytes.NewReader(streamErrorBody))} - s.handleFailoverSideEffects(ctx, syntheticResp, account, reqModel) - } - failure := classification.FailoverError(streamErrorBody) - if activeRecoveryDiagnostic != nil { - failure.NextAccountAction = NextAccountStop - failure.SuppressAccountHealthPenalty = true - } - return nil, failure - } - if activeRecoveryDiagnostic != nil { - activeRecoveryDiagnostic.Outcome, activeRecoveryDiagnostic.Reason = "response_read_failed", "stream_incomplete" - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - ProxyID: opsUpstreamProxyID(account), ProxyName: opsUpstreamProxyName(account), - Platform: account.Platform, AccountID: account.ID, AccountName: account.Name, - UpstreamStatusCode: http.StatusBadGateway, UpstreamHTTPStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "stream_failure", Message: sanitizeStreamError(err), - }) - var failure *UpstreamFailoverError - if errors.As(err, &failure) { - failure.NextAccountAction = NextAccountStop - failure.SuppressAccountHealthPenalty = true - } - } - // 流中断（缺失 terminal 事件、读错误、数据间隔超时等）时保留已观测到的 - // usage 与错误一起返回，handler 在错误处理完成后照常提交 usage 记录。 - if partial := partialStreamUsageResult(c, resp, streamResult, originalModel, mappedModel, startTime, err); partial != nil { - return partial, err - } - return nil, err + writerSizeBeforeStream := c.Writer.Size() + streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, reqModel, shouldMimicClaudeCode) + if err != nil { + var sseErr *sseStreamErrorEventError + if errors.As(err, &sseErr) { + return nil, s.sseStreamErrorEventFailover(ctx, c, account, resp, sseErr.RawData, reqModel, c.Writer.Size() != writerSizeBeforeStream) } - usage = streamResult.usage - firstTokenMs = streamResult.firstTokenMs - clientDisconnect = streamResult.clientDisconnect - break + // 流中断（缺失 terminal 事件、读错误、数据间隔超时等）时保留已观测到的 + // usage 与错误一起返回，handler 在错误处理完成后照常提交 usage 记录。 + if partial := partialStreamUsageResult(c, resp, streamResult, originalModel, mappedModel, startTime, err); partial != nil { + return partial, err + } + return nil, err } + usage = streamResult.usage + firstTokenMs = streamResult.firstTokenMs + clientDisconnect = streamResult.clientDisconnect } else { usage, err = s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, reqModel) if err != nil { @@ -1166,3 +1049,57 @@ func (s *GatewayService) isStickyAccountUpstreamRestricted(ctx context.Context, } return s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel) } + +// sseStreamErrorEventFailover 把上游 HTTP 200 + SSE 流体内的 event:error 帧转为 +// UpstreamFailoverError：尚未向客户端写出内容且错误类型为 overloaded_error 时按 529 +// 处理并触发账号侧副作用，其余按 403；同时登记 stream_error 运维记录。 +func (s *GatewayService) sseStreamErrorEventFailover(ctx context.Context, c *gin.Context, account *Account, resp *http.Response, rawData, model string, outputStarted bool) *UpstreamFailoverError { + body := []byte(rawData) + semanticStatus := http.StatusForbidden + if !outputStarted && gjson.GetBytes(body, "error.type").String() == "overloaded_error" { + semanticStatus = 529 + syntheticResp := &http.Response{ + StatusCode: semanticStatus, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(body)), + } + s.handleFailoverSideEffects(ctx, syntheticResp, account, model) + } + + upstreamMsg := sanitizeUpstreamErrorMessage( + strings.TrimSpace(extractUpstreamErrorMessage(body)), + ) + + upstreamDetail := "" + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes + if maxBytes <= 0 { + maxBytes = 2048 + } + upstreamDetail = truncateString(rawData, maxBytes) + } + + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + ProxyID: opsUpstreamProxyID(account), + ProxyName: opsUpstreamProxyName(account), + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: semanticStatus, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "stream_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + logger.LegacyPrintf("service.gateway", + "[Forward] SSE error event in stream: Account=%d(%s) RequestID=%s Body=%s", + account.ID, account.Name, resp.Header.Get("x-request-id"), + truncateString(rawData, 1000), + ) + + return &UpstreamFailoverError{ + StatusCode: semanticStatus, + ResponseBody: body, + } +} diff --git a/backend/internal/service/gateway_forward_as_chat_completions.go b/backend/internal/service/gateway_forward_as_chat_completions.go index 108457c2a5f7..2d7a36f78ef1 100644 --- a/backend/internal/service/gateway_forward_as_chat_completions.go +++ b/backend/internal/service/gateway_forward_as_chat_completions.go @@ -188,7 +188,15 @@ func (s *GatewayService) forwardAsChatCompletions( var result *ForwardResult var handleErr error if clientStream { + writerSizeBeforeStream := c.Writer.Size() result, handleErr = s.handleCCStreamingFromAnthropic(resp, c, originalModel, mappedModel, reasoningEffort, startTime) + var sseErr *sseStreamErrorEventError + switch { + case errors.As(handleErr, &sseErr): + return nil, s.sseStreamErrorEventFailover(ctx, c, account, resp, sseErr.RawData, mappedModel, c.Writer.Size() != writerSizeBeforeStream) + case errors.Is(handleErr, errAnthropicNativeStreamIdle) && s.rateLimitService != nil: + s.rateLimitService.HandleStreamTimeout(ctx, account, originalModel) + } } else { result, handleErr = s.handleCCBufferedFromAnthropic(resp, c, originalModel, mappedModel, reasoningEffort, startTime) } @@ -382,26 +390,23 @@ func (s *GatewayService) handleCCStreamingFromAnthropic( var usage ClaudeUsage var firstTokenMs *int firstChunk := true - var protocolErr error + writerSizeBeforeStream := c.Writer.Size() - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.cfg.Gateway.MaxLineSize - } - scanner.Buffer(make([]byte, 0, 64*1024), maxLineSize) + stream := newAnthropicCompatStream(s.cfg, c, resp, "forward_as_cc", requestID) + defer stream.stop() resultWithUsage := func() *ForwardResult { return &ForwardResult{ - RequestID: requestID, - UpstreamHeaders: resp.Header, - Usage: usage, - Model: originalModel, - UpstreamModel: mappedModel, - ReasoningEffort: reasoningEffort, - Stream: true, - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, + RequestID: requestID, + UpstreamHeaders: resp.Header, + Usage: usage, + Model: originalModel, + UpstreamModel: mappedModel, + ReasoningEffort: reasoningEffort, + Stream: true, + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ClientDisconnect: stream.clientDisconnected, } } @@ -419,36 +424,20 @@ func (s *GatewayService) handleCCStreamingFromAnthropic( return false } - writeResponseEvent := func(event *apicompat.ResponsesStreamEvent) bool { - if protocolErr != nil { - return false - } - chunks := apicompat.ResponsesEventToChatChunks(event, ccState) - if ccState.ProtocolError != "" { - protocolErr = errors.New(ccState.ProtocolError) - MarkResponseCommitted(c) - _, _ = fmt.Fprint(c.Writer, buildChatStreamErrorSSE("upstream_protocol_error", ccState.ProtocolError)+"data: [DONE]\n\n") - c.Writer.Flush() - return false // Keep reading the upstream usage, without further output. - } - for _, chunk := range chunks { - if writeChunk(chunk) { - return true - } - } - return false - } - - processAnthropicEvent := func(event *apicompat.AnthropicStreamEvent) bool { + processAnthropicEvent := func(event *apicompat.AnthropicStreamEvent) { if event == nil { - return false + return } if event.Type == "ping" { + if stream.clientDisconnected { + return + } if _, err := fmt.Fprint(c.Writer, ": ping\n\n"); err != nil { - return true + stream.markClientDisconnected() + return } c.Writer.Flush() - return false + return } if firstChunk { firstChunk = false @@ -464,8 +453,8 @@ func (s *GatewayService) handleCCStreamingFromAnthropic( if event.Type == "message_start" && event.Message != nil { mergeAnthropicUsage(&usage, event.Message.Usage) } - if protocolErr != nil { - return false + if stream.clientDisconnected { + return } // Keep the outward Responses/Chat usage on the same normalized buckets used @@ -476,25 +465,32 @@ func (s *GatewayService) handleCCStreamingFromAnthropic( // Chain: Anthropic event → Responses events → CC chunks responsesEvents := apicompat.AnthropicEventToResponsesEvents(event, anthState) for _, resEvt := range responsesEvents { - if writeResponseEvent(&resEvt) { - return true + ccChunks := apicompat.ResponsesEventToChatChunks(&resEvt, ccState) + for _, chunk := range ccChunks { + if disconnected := writeChunk(chunk); disconnected { + stream.markClientDisconnected() + return + } } } c.Writer.Flush() - return false } - for scanner.Scan() { - line := scanner.Text() + for { + line, ok := stream.nextLine() + if !ok { + break + } // 与缓冲路径一致：接受 SSE 紧凑格式（冒号后无空格，#4653 同根因）。 if _, ok := extractOpenAISSEEventLine(line); !ok { continue } - if !scanner.Scan() { + dataLine, ok := stream.nextLine() + if !ok { break } - payload, ok := extractOpenAISSEDataLine(scanner.Text()) + payload, ok := extractOpenAISSEDataLine(dataLine) if !ok { continue } @@ -508,30 +504,30 @@ func (s *GatewayService) handleCCStreamingFromAnthropic( // The intermediate Responses converter synthesizes usage even when absent. ccState.IncludeUsage = ccState.IncludeUsage || anthropicChatStreamHasUsage(&event, payload) - if processAnthropicEvent(&event) { - return resultWithUsage(), protocolErr + stream.observe(&event, payload) + if stream.upstreamErrType != "" { + break } + processAnthropicEvent(&event) } - if err := scanner.Err(); err != nil { - if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { - logger.L().Warn("forward_as_cc stream: read error", - zap.Error(err), - zap.String("request_id", requestID), - ) - } + if failure := stream.failure(); failure != nil { + return finishAnthropicCompatStreamFailure(c, writerSizeBeforeStream, failure, resultWithUsage(), func() { + writeAnthropicCompatCCStreamFailure(c, failure) + }) + } + if stream.clientDisconnected { + return resultWithUsage(), nil } // Finalize both state machines finalResEvents := apicompat.FinalizeAnthropicResponsesStream(anthState) for _, resEvt := range finalResEvents { - if writeResponseEvent(&resEvt) { - return resultWithUsage(), protocolErr + ccChunks := apicompat.ResponsesEventToChatChunks(&resEvt, ccState) + for _, chunk := range ccChunks { + writeChunk(chunk) //nolint:errcheck } } - if protocolErr != nil { - return resultWithUsage(), protocolErr - } finalCCChunks := apicompat.FinalizeResponsesChatStream(ccState) for _, chunk := range finalCCChunks { writeChunk(chunk) //nolint:errcheck diff --git a/backend/internal/service/gateway_forward_as_chat_completions_ping_test.go b/backend/internal/service/gateway_forward_as_chat_completions_ping_test.go index df08d20f78ce..7eda0f90d9f4 100644 --- a/backend/internal/service/gateway_forward_as_chat_completions_ping_test.go +++ b/backend/internal/service/gateway_forward_as_chat_completions_ping_test.go @@ -190,15 +190,16 @@ func TestHandleCCStreamingFromAnthropic_PingWriteFailureStopsDownstream(t *testi ) require.NoError(t, err) require.True(t, writer.failed) + require.True(t, result.ClientDisconnect) require.Zero(t, writer.writesAfterError) require.Zero(t, writer.flushesAfterError) require.NotContains(t, writer.Body.String(), "[DONE]") - require.Zero(t, result.Usage.OutputTokens) + require.Equal(t, 7, result.Usage.OutputTokens) + require.NotNil(t, result.FirstTokenMs) if withUsage { require.Equal(t, 10, result.Usage.InputTokens) } else { require.Zero(t, result.Usage.InputTokens) - require.Nil(t, result.FirstTokenMs) } }) } diff --git a/backend/internal/service/gateway_forward_as_responses.go b/backend/internal/service/gateway_forward_as_responses.go index 143ee21d8036..3eeac70f7ccf 100644 --- a/backend/internal/service/gateway_forward_as_responses.go +++ b/backend/internal/service/gateway_forward_as_responses.go @@ -199,7 +199,15 @@ func (s *GatewayService) forwardAsResponses( var result *ForwardResult var handleErr error if clientStream { + writerSizeBeforeStream := c.Writer.Size() result, handleErr = s.handleResponsesStreamingResponse(resp, c, originalModel, mappedModel, reasoningEffort, startTime, clientToolMapping) + var sseErr *sseStreamErrorEventError + switch { + case errors.As(handleErr, &sseErr): + return nil, s.sseStreamErrorEventFailover(ctx, c, account, resp, sseErr.RawData, mappedModel, c.Writer.Size() != writerSizeBeforeStream) + case errors.Is(handleErr, errAnthropicNativeStreamIdle) && s.rateLimitService != nil: + s.rateLimitService.HandleStreamTimeout(ctx, account, originalModel) + } } else { result, handleErr = s.handleResponsesBufferedStreamingResponse(resp, c, originalModel, mappedModel, reasoningEffort, startTime, clientToolMapping) } @@ -546,30 +554,28 @@ func (s *GatewayService) handleResponsesStreamingResponse( var usage ClaudeUsage var firstTokenMs *int firstChunk := true + writerSizeBeforeStream := c.Writer.Size() - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.cfg.Gateway.MaxLineSize - } - scanner.Buffer(make([]byte, 0, 64*1024), maxLineSize) + stream := newAnthropicCompatStream(s.cfg, c, resp, "forward_as_responses", requestID) + defer stream.stop() resultWithUsage := func() *ForwardResult { return &ForwardResult{ - RequestID: requestID, - UpstreamHeaders: resp.Header, - Usage: usage, - Model: originalModel, - UpstreamModel: mappedModel, - ReasoningEffort: reasoningEffort, - Stream: true, - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, + RequestID: requestID, + UpstreamHeaders: resp.Header, + Usage: usage, + Model: originalModel, + UpstreamModel: mappedModel, + ReasoningEffort: reasoningEffort, + Stream: true, + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ClientDisconnect: stream.clientDisconnected, } } // processEvent handles a single parsed Anthropic SSE event. - processEvent := func(event *apicompat.AnthropicStreamEvent) bool { + processEvent := func(event *apicompat.AnthropicStreamEvent) { if firstChunk { firstChunk = false ms := int(time.Since(startTime).Milliseconds()) @@ -584,6 +590,9 @@ func (s *GatewayService) handleResponsesStreamingResponse( if event.Type == "message_start" && event.Message != nil { mergeAnthropicUsage(&usage, event.Message.Usage) } + if stream.clientDisconnected { + return + } // Keep the terminal Responses usage aligned with the normalized billing // buckets. Normalize the converter input too, so message handlers cannot @@ -614,20 +623,25 @@ func (s *GatewayService) handleResponsesStreamingResponse( for _, restored := range payloads { eventType := gjson.GetBytes(restored, "type").String() if _, err := fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", eventType, restored); err != nil { - logger.L().Info("forward_as_responses stream: client disconnected", - zap.String("request_id", requestID), - ) - return true // client disconnected + stream.markClientDisconnected() + return } } } if len(events) > 0 { c.Writer.Flush() } - return false } finalizeStream := func() (*ForwardResult, error) { + if failure := stream.failure(); failure != nil { + return finishAnthropicCompatStreamFailure(c, writerSizeBeforeStream, failure, resultWithUsage(), func() { + writeAnthropicCompatResponsesStreamFailure(c, state, clientToolRestorer, failure) + }) + } + if stream.clientDisconnected { + return resultWithUsage(), nil + } if finalEvents := apicompat.FinalizeAnthropicResponsesStream(state); len(finalEvents) > 0 { for _, evt := range finalEvents { sse, err := apicompat.ResponsesEventToSSE(evt) @@ -643,18 +657,21 @@ func (s *GatewayService) handleResponsesStreamingResponse( } // Read Anthropic SSE events - for scanner.Scan() { - line := scanner.Text() + for { + line, ok := stream.nextLine() + if !ok { + break + } eventType, ok := parseAnthropicSSEField(line, "event") if !ok { continue } // Read data line - if !scanner.Scan() { + dataLine, ok := stream.nextLine() + if !ok { break } - dataLine := scanner.Text() payload, ok := parseAnthropicSSEField(dataLine, "data") if !ok { continue @@ -670,18 +687,11 @@ func (s *GatewayService) handleResponsesStreamingResponse( continue } - if processEvent(&event) { - return resultWithUsage(), nil - } - } - - if err := scanner.Err(); err != nil { - if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { - logger.L().Warn("forward_as_responses stream: read error", - zap.Error(err), - zap.String("request_id", requestID), - ) + stream.observe(&event, payload) + if stream.upstreamErrType != "" { + break } + processEvent(&event) } return finalizeStream() diff --git a/backend/internal/service/gateway_forward_compat_client_disconnect_test.go b/backend/internal/service/gateway_forward_compat_client_disconnect_test.go new file mode 100644 index 000000000000..d41e42a4ada4 --- /dev/null +++ b/backend/internal/service/gateway_forward_compat_client_disconnect_test.go @@ -0,0 +1,831 @@ +//go:build unit + +package service + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +// disconnectOnMarkerResponseWriter 在首次写入包含 marker 的数据时返回错误，模拟客户端中途断开。 +type disconnectOnMarkerResponseWriter struct { + *httptest.ResponseRecorder + marker string + failed bool + writesAfterError int + flushesAfterError int +} + +func (writer *disconnectOnMarkerResponseWriter) Write(data []byte) (int, error) { + if writer.failed { + writer.writesAfterError++ + return 0, errors.New("client disconnected") + } + if strings.Contains(string(data), writer.marker) { + writer.failed = true + return 0, errors.New("client disconnected") + } + return writer.ResponseRecorder.Write(data) +} + +func (writer *disconnectOnMarkerResponseWriter) WriteString(data string) (int, error) { + return writer.Write([]byte(data)) +} + +func (writer *disconnectOnMarkerResponseWriter) Flush() { + if writer.failed { + writer.flushesAfterError++ + } + writer.ResponseRecorder.Flush() +} + +func anthropicStreamWithMidStreamText() string { + return strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg_disc","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","usage":{"input_tokens":20,"cache_read_input_tokens":11,"cache_creation_input_tokens":4,"output_tokens":1}}}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"first"}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"second"}}`, + ``, + `event: ping`, + `data: {"type":"ping"}`, + ``, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":0}`, + ``, + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":42}}`, + ``, + `event: message_stop`, + `data: {"type":"message_stop"}`, + ``, + }, "\n") +} + +// stallingAnthropicUpstream 写完 stream 后保持连接不关闭，模拟上游挂住。 +func stallingAnthropicUpstream(t *testing.T, stream string) (*http.Response, *io.PipeWriter) { + t.Helper() + pr, pw := io.Pipe() + go func() { + _, _ = io.WriteString(pw, stream) + }() + t.Cleanup(func() { + _ = pw.Close() + _ = pr.Close() + }) + return &http.Response{Header: http.Header{}, Body: pr}, pw +} + +func newCompatStreamIdleTestService(intervalSec int) *GatewayService { + return &GatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{StreamDataIntervalTimeout: intervalSec}}} +} + +type compatStreamHandler func(*GatewayService, *http.Response, *gin.Context) (*ForwardResult, error) + +func compatStreamHandlers() map[string]compatStreamHandler { + return map[string]compatStreamHandler{ + "chat_completions": func(svc *GatewayService, resp *http.Response, c *gin.Context) (*ForwardResult, error) { + return svc.handleCCStreamingFromAnthropic(resp, c, "claude-sonnet-4-6", "claude-sonnet-4-6", nil, time.Now()) + }, + "responses": func(svc *GatewayService, resp *http.Response, c *gin.Context) (*ForwardResult, error) { + return svc.handleResponsesStreamingResponse(resp, c, "claude-sonnet-4-6", "claude-sonnet-4-6", nil, time.Now(), apicompat.ResponsesClientToolMapping{}) + }, + } +} + +// compatStreamTerminalMarkers 是两条兼容流按成功收尾时写出的完成标记。 +var compatStreamTerminalMarkers = map[string]string{ + "chat_completions": "[DONE]", + "responses": "response.completed", +} + +// compatStreamFailureMarkers 是两条兼容流写出终止错误时的标记。 +var compatStreamFailureMarkers = map[string]string{ + "chat_completions": `data: {"error":`, + "responses": "event: response.failed", +} + +// anthropicStreamWithoutMessageStop 以带 stop_reason 的 message_delta 结束，不含 message_stop。 +func anthropicStreamWithoutMessageStop() string { + return strings.TrimSuffix(anthropicStreamWithMidStreamText(), "event: message_stop\ndata: {\"type\":\"message_stop\"}\n") +} + +// anthropicStreamTruncatedBeforeMessageDelta 在 message_delta 之前中断，usage 只有 message_start 的部分。 +func anthropicStreamTruncatedBeforeMessageDelta() string { + full := anthropicStreamWithMidStreamText() + return full[:strings.Index(full, "event: message_delta")] +} + +type compatStreamOutcome struct { + result *ForwardResult + err error +} + +func runCompatStreamWithDeadline(t *testing.T, handle compatStreamHandler, svc *GatewayService, resp *http.Response, c *gin.Context) compatStreamOutcome { + t.Helper() + done := make(chan compatStreamOutcome, 1) + go func() { + result, err := handle(svc, resp, c) + done <- compatStreamOutcome{result: result, err: err} + }() + select { + case got := <-done: + return got + case <-time.After(5 * time.Second): + t.Fatal("stream did not end after upstream idle interval") + return compatStreamOutcome{} + } +} + +func TestAnthropicCompatStreaming_IdleUpstreamEndsWithinInterval(t *testing.T) { + for name, handle := range compatStreamHandlers() { + for _, disconnect := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/client_disconnected=%t", name, disconnect), func(t *testing.T) { + gin.SetMode(gin.TestMode) + marker := "\x00never-written" + if disconnect { + marker = "second" + } + writer := &disconnectOnMarkerResponseWriter{ResponseRecorder: httptest.NewRecorder(), marker: marker} + c, _ := gin.CreateTestContext(writer) + resp, pw := stallingAnthropicUpstream(t, anthropicStreamTruncatedBeforeMessageDelta()) + + got := runCompatStreamWithDeadline(t, handle, newCompatStreamIdleTestService(1), resp, c) + + if disconnect { + require.NoError(t, got.err) + } else { + require.ErrorIs(t, got.err, errAnthropicNativeStreamIdle) + body := writer.Body.String() + require.NotContains(t, body, compatStreamTerminalMarkers[name], "a stream cut by idle timeout must not be reported as completed") + require.Contains(t, body, compatStreamFailureMarkers[name]) + require.Contains(t, body, "stream_timeout") + require.True(t, IsResponseCommitted(c)) + } + require.NotNil(t, got.result) + require.Equal(t, disconnect, got.result.ClientDisconnect) + require.Equal(t, disconnect, writer.failed) + require.Equal(t, 20, got.result.Usage.InputTokens) + require.Equal(t, 1, got.result.Usage.OutputTokens) + require.Zero(t, writer.writesAfterError) + _, werr := io.WriteString(pw, "event: ping\n") + require.Error(t, werr, "upstream body must be closed after idle timeout") + }) + } + } +} + +func TestAnthropicCompatStreaming_IdleAfterTerminalEventCompletes(t *testing.T) { + for name, handle := range compatStreamHandlers() { + for streamName, stream := range map[string]string{ + "message_stop": anthropicStreamWithMidStreamText(), + "stop_reason": anthropicStreamWithoutMessageStop(), + } { + t.Run(name+"/"+streamName, func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + resp, _ := stallingAnthropicUpstream(t, stream) + + got := runCompatStreamWithDeadline(t, handle, newCompatStreamIdleTestService(1), resp, c) + + require.NoError(t, got.err) + require.NotNil(t, got.result) + require.Equal(t, 42, got.result.Usage.OutputTokens) + require.Contains(t, rec.Body.String(), compatStreamTerminalMarkers[name]) + require.NotContains(t, rec.Body.String(), compatStreamFailureMarkers[name]) + }) + } + } +} + +// errAfterReader 读完 data 后返回 err，模拟上游连接中途断开。 +type errAfterReader struct { + data *strings.Reader + err error +} + +func (r *errAfterReader) Read(p []byte) (int, error) { + if r.data.Len() > 0 { + return r.data.Read(p) + } + return 0, r.err +} + +func TestAnthropicCompatStreaming_ReadErrorBeforeTerminalEventFails(t *testing.T) { + messageStartOnly := strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg_disc","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","usage":{"input_tokens":20,"output_tokens":1}}}`, + ``, + }, "\n") + "\n" + for name, handle := range compatStreamHandlers() { + for _, tc := range []struct { + name string + stream string + wantErr bool + wantResult bool + }{ + {name: "after_text", stream: anthropicStreamTruncatedBeforeMessageDelta(), wantErr: true, wantResult: true}, + {name: "after_message_start", stream: messageStartOnly, wantErr: true, wantResult: true}, + {name: "before_any_event", stream: "", wantErr: true}, + {name: "after_stop_reason", stream: anthropicStreamWithoutMessageStop(), wantResult: true}, + {name: "after_message_stop_without_stop_reason", stream: strings.Replace(anthropicStreamWithMidStreamText(), `"delta":{"stop_reason":"end_turn"}`, `"delta":{}`, 1), wantResult: true}, + } { + t.Run(name+"/"+tc.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := io.NopCloser(&errAfterReader{data: strings.NewReader(tc.stream), err: io.ErrUnexpectedEOF}) + + result, err := handle(&GatewayService{}, &http.Response{Header: http.Header{}, Body: body}, c) + + if tc.wantErr { + require.ErrorIs(t, err, io.ErrUnexpectedEOF) + require.NotErrorIs(t, err, errAnthropicNativeStreamIdle) + require.NotContains(t, rec.Body.String(), compatStreamTerminalMarkers[name]) + require.Contains(t, rec.Body.String(), compatStreamFailureMarkers[name]) + require.Contains(t, rec.Body.String(), "stream_read_error") + } else { + require.NoError(t, err) + require.Contains(t, rec.Body.String(), compatStreamTerminalMarkers[name]) + } + if !tc.wantResult { + require.Nil(t, result) + return + } + require.NotNil(t, result) + require.False(t, result.ClientDisconnect) + require.Equal(t, 20, result.Usage.InputTokens) + }) + } + } +} + +// compatResponsesSSEEvents 解析 Responses SSE 输出中各 data 帧的 type 与 sequence_number。 +func compatResponsesSSEEvents(t *testing.T, body string) []gjson.Result { + t.Helper() + var events []gjson.Result + for _, line := range strings.Split(body, "\n") { + if payload, ok := strings.CutPrefix(line, "data: "); ok { + require.True(t, gjson.Valid(payload), payload) + events = append(events, gjson.Parse(payload)) + } + } + return events +} + +func TestAnthropicCompatStreaming_EOFWithoutTerminalEventFails(t *testing.T) { + truncatedToolUse := strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg_disc","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","usage":{"input_tokens":20,"output_tokens":1}}}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_1","name":"write_file","input":{}}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"path\":\"/tmp/a\",\"content\":\"hal"}}`, + ``, + }, "\n") + "\n" + for name, handle := range compatStreamHandlers() { + for _, tc := range []struct { + name string + stream string + wantErrText string + wantType string + }{ + {name: "truncated", stream: anthropicStreamTruncatedBeforeMessageDelta(), wantErrText: "missing terminal event", wantType: "stream_incomplete"}, + {name: "truncated_tool_use", stream: truncatedToolUse, wantErrText: "missing terminal event", wantType: "stream_incomplete"}, + } { + t.Run(name+"/"+tc.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + + result, err := handle(&GatewayService{}, &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader(tc.stream))}, c) + + require.ErrorContains(t, err, tc.wantErrText) + require.NotNil(t, result) + require.Equal(t, 20, result.Usage.InputTokens) + require.True(t, IsResponseCommitted(c)) + body := rec.Body.String() + require.NotContains(t, body, compatStreamTerminalMarkers[name]) + require.NotContains(t, body, `"finish_reason":"`) + require.Equal(t, 1, strings.Count(body, compatStreamFailureMarkers[name])) + if name != "responses" { + require.Contains(t, body, tc.wantType) + return + } + events := compatResponsesSSEEvents(t, body) + require.GreaterOrEqual(t, len(events), 2) + failed := events[len(events)-1] + require.Equal(t, "response.failed", failed.Get("type").String()) + require.Equal(t, "msg_disc", failed.Get("response.id").String()) + require.Equal(t, "failed", failed.Get("response.status").String()) + require.Equal(t, tc.wantType, failed.Get("response.error.code").String()) + require.Equal(t, events[len(events)-2].Get("sequence_number").Int()+1, failed.Get("sequence_number").Int()) + }) + } + } +} + +func TestAnthropicCompatStreaming_EOFBeforeMessageStartKeepsExistingEnding(t *testing.T) { + for name, handle := range compatStreamHandlers() { + t.Run(name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + + result, err := handle(&GatewayService{}, &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader("event: ping\ndata: {\"type\":\"ping\"}\n\n"))}, c) + + require.NoError(t, err) + require.NotNil(t, result) + require.NotContains(t, rec.Body.String(), compatStreamFailureMarkers[name]) + }) + } +} + +func TestAnthropicCompatStreaming_DoneLineCountsAsTerminalEvent(t *testing.T) { + for name, handle := range compatStreamHandlers() { + t.Run(name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + stream := anthropicStreamTruncatedBeforeMessageDelta() + "data: [DONE]\n\n" + + result, err := handle(&GatewayService{}, &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader(stream))}, c) + + require.NoError(t, err) + require.NotNil(t, result) + require.NotContains(t, rec.Body.String(), compatStreamFailureMarkers[name]) + }) + } +} + +func anthropicStreamWithInBandError(messageBeforeError, after string) string { + return strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg_disc","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","usage":{"input_tokens":20,"output_tokens":1}}}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"second"}}`, + ``, + `event: error`, + `data: {"type":"error","error":{"type":"overloaded_error","message":` + strconv.Quote(messageBeforeError) + `}}`, + ``, + }, "\n") + "\n" + after +} + +func TestAnthropicCompatStreaming_UpstreamErrorEventEndsStreamWithoutWaiting(t *testing.T) { + for name, handle := range compatStreamHandlers() { + for _, disconnect := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/client_disconnected=%t", name, disconnect), func(t *testing.T) { + gin.SetMode(gin.TestMode) + marker := "\x00never-written" + if disconnect { + marker = "second" + } + writer := &disconnectOnMarkerResponseWriter{ResponseRecorder: httptest.NewRecorder(), marker: marker} + c, _ := gin.CreateTestContext(writer) + resp, pw := stallingAnthropicUpstream(t, anthropicStreamWithInBandError("Overloaded", "")) + + got := runCompatStreamWithDeadline(t, handle, newCompatStreamIdleTestService(0), resp, c) + + if disconnect { + require.NoError(t, got.err) + require.NotNil(t, got.result) + require.True(t, got.result.ClientDisconnect) + require.Equal(t, 20, got.result.Usage.InputTokens) + require.Zero(t, writer.writesAfterError) + } else { + var sseErr *sseStreamErrorEventError + require.ErrorAs(t, got.err, &sseErr) + require.Nil(t, got.result, "an upstream error event is not billed while the client is online") + require.Equal(t, 1, strings.Count(writer.Body.String(), compatStreamFailureMarkers[name])) + require.Contains(t, writer.Body.String(), "Overloaded") + } + _, werr := io.WriteString(pw, "event: ping\n") + require.Error(t, werr, "upstream body must be closed after an upstream error event") + }) + } + } +} + +func TestAnthropicCompatStreaming_UpstreamErrorEventStopsReadingBufferedEvents(t *testing.T) { + after := strings.Join([]string{ + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"LEAKED"}}`, + ``, + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":9}}`, + ``, + `event: message_stop`, + `data: {"type":"message_stop"}`, + ``, + }, "\n") + "\n" + for name, handle := range compatStreamHandlers() { + t.Run(name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + + result, err := handle(&GatewayService{}, &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader(anthropicStreamWithInBandError("Overloaded", after)))}, c) + + var sseErr *sseStreamErrorEventError + require.ErrorAs(t, err, &sseErr) + require.Nil(t, result) + body := rec.Body.String() + require.NotContains(t, body, "LEAKED") + require.NotContains(t, body, compatStreamTerminalMarkers[name]) + require.Equal(t, 1, strings.Count(body, compatStreamFailureMarkers[name])) + }) + } +} + +func TestAnthropicCompatStreaming_UpstreamErrorEventMessageIsSanitized(t *testing.T) { + for name, handle := range compatStreamHandlers() { + t.Run(name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + stream := anthropicStreamWithInBandError("upstream https://example.test/v1?key=SECRET failed", "") + + _, err := handle(&GatewayService{}, &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader(stream))}, c) + + require.Error(t, err) + require.Contains(t, rec.Body.String(), "key=***") + require.NotContains(t, rec.Body.String(), "SECRET") + }) + } +} + +func TestAnthropicCompatStreaming_UpstreamErrorEventBeforeOutputWritesNothing(t *testing.T) { + stream := "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n" + for name, handle := range compatStreamHandlers() { + t.Run(name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + + result, err := handle(&GatewayService{}, &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader(stream))}, c) + + var sseErr *sseStreamErrorEventError + require.ErrorAs(t, err, &sseErr) + require.Nil(t, result) + require.Zero(t, rec.Body.Len(), "the handler fails over to another account") + require.False(t, IsResponseCommitted(c)) + }) + } +} + +func TestForwardAsCompatStreaming_UpstreamErrorEventFailsOver(t *testing.T) { + beforeOutput := "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n" + for _, endpoint := range []string{"chat_completions", "responses"} { + for _, tc := range []struct { + name string + stream string + wantStatus int + }{ + {name: "before_output", stream: beforeOutput, wantStatus: 529}, + {name: "after_output", stream: anthropicStreamWithInBandError("Overloaded", ""), wantStatus: http.StatusForbidden}, + } { + t.Run(endpoint+"/"+tc.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := `{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi"}],"stream":true}` + path := "/v1/chat/completions" + if endpoint == "responses" { + body = `{"model":"claude-sonnet-4-6","input":"hi","stream":true}` + path = "/v1/responses" + } + c.Request = httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)) + resp := &http.Response{StatusCode: http.StatusOK, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(tc.stream))} + svc := &GatewayService{ + cfg: rawChatCompletionsTestConfig(), + httpUpstream: &anthropicHTTPUpstreamRecorder{resp: resp}, + rateLimitService: &RateLimitService{accountRepo: &errorPolicyRepoStub{}, cfg: &config.Config{}}, + } + account := &Account{ID: 1, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "sk-test"}} + + var result *ForwardResult + var err error + if endpoint == "responses" { + result, err = svc.ForwardAsResponses(context.Background(), c, account, []byte(body), nil) + } else { + result, err = svc.ForwardAsChatCompletions(context.Background(), c, account, []byte(body), nil) + } + + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, tc.wantStatus, failoverErr.StatusCode) + require.Nil(t, result) + }) + } + } +} + +func TestAnthropicCompatStreaming_LineTooLongFails(t *testing.T) { + for name, handle := range compatStreamHandlers() { + t.Run(name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + svc := &GatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: 128 * 1024}}} + stream := anthropicStreamTruncatedBeforeMessageDelta() + "event: content_block_delta\ndata: " + strings.Repeat("x", 256*1024) + "\n\n" + + result, err := handle(svc, &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader(stream))}, c) + + require.ErrorIs(t, err, bufio.ErrTooLong) + require.NotNil(t, result) + require.Contains(t, rec.Body.String(), compatStreamFailureMarkers[name]) + require.Contains(t, rec.Body.String(), "response_too_large") + }) + } +} + +func TestHandleResponsesStreamingResponse_FailureSequenceFollowsRestoredEvents(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + stream := strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg_tool","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-6","usage":{"input_tokens":20,"output_tokens":1}}}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_1","name":"apply_patch","input":{}}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"input\":"}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"*** Begin"}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":" Patch"}}`, + ``, + }, "\n") + "\n" + mapping := apicompat.ResponsesClientToolMapping{CustomTools: map[string]bool{"apply_patch": true}} + + _, err := (&GatewayService{}).handleResponsesStreamingResponse( + &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader(stream))}, + c, "claude-sonnet-4-6", "claude-sonnet-4-6", nil, time.Now(), mapping, + ) + + require.ErrorContains(t, err, "missing terminal event") + events := compatResponsesSSEEvents(t, rec.Body.String()) + require.NotEmpty(t, events) + require.Equal(t, "response.failed", events[len(events)-1].Get("type").String()) + for i, event := range events { + require.Equal(t, int64(i), event.Get("sequence_number").Int(), "sequence_number must be contiguous: %s", event.Raw) + } +} + +func TestAnthropicCompatStream_ClientGoneAtStreamEndMarksDisconnected(t *testing.T) { + gin.SetMode(gin.TestMode) + // 泵 goroutine 结束与请求 ctx 取消同时就绪时 select 随机选择，多轮重复以覆盖两种选择。 + for round := 0; round < 64; round++ { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx) + stream := newAnthropicCompatStream(nil, c, &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader(""))}, "test", "req") + time.Sleep(time.Millisecond) + + _, ok := stream.nextLine() + + require.False(t, ok) + require.True(t, stream.clientDisconnected, "round %d", round) + require.Nil(t, stream.failure()) + stream.stop() + } +} + +func TestAnthropicCompatStreaming_RequestContextCanceledDrainsWithoutWriting(t *testing.T) { + for name, handle := range compatStreamHandlers() { + for _, delayed := range []bool{true, false} { + t.Run(fmt.Sprintf("%s/upstream_delayed=%t", name, delayed), func(t *testing.T) { + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx) + var body io.ReadCloser = io.NopCloser(strings.NewReader(anthropicStreamWithMidStreamText())) + if delayed { + pr, pw := io.Pipe() + t.Cleanup(func() { _ = pr.Close() }) + go func() { + time.Sleep(20 * time.Millisecond) + _, _ = io.WriteString(pw, anthropicStreamWithMidStreamText()) + _ = pw.Close() + }() + body = pr + } + + got := runCompatStreamWithDeadline(t, handle, newCompatStreamIdleTestService(0), &http.Response{Header: http.Header{}, Body: body}, c) + + require.NoError(t, got.err) + require.NotNil(t, got.result) + require.True(t, got.result.ClientDisconnect) + require.Equal(t, 20, got.result.Usage.InputTokens) + require.Equal(t, 42, got.result.Usage.OutputTokens) + require.NotContains(t, rec.Body.String(), compatStreamTerminalMarkers[name]) + }) + } + } +} + +func TestAnthropicCompatStream_ClientDisconnectBoundsDrainWhenIntervalDisabled(t *testing.T) { + for _, tc := range []struct { + name string + intervalSec int + wantInterval time.Duration + }{ + {name: "interval_disabled", intervalSec: 0, wantInterval: anthropicCompatDisconnectDrainInterval}, + {name: "interval_configured", intervalSec: 30, wantInterval: 30 * time.Second}, + } { + t.Run(tc.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + pr, pw := io.Pipe() + t.Cleanup(func() { _ = pw.Close(); _ = pr.Close() }) + cfg := &config.Config{Gateway: config.GatewayConfig{StreamDataIntervalTimeout: tc.intervalSec}} + stream := newAnthropicCompatStream(cfg, c, &http.Response{Header: http.Header{}, Body: pr}, "test", "req") + defer stream.stop() + require.Equal(t, tc.intervalSec > 0, stream.pump.timer != nil) + + stream.markClientDisconnected() + + require.True(t, stream.clientDisconnected) + require.NotNil(t, stream.pump.timer) + require.Equal(t, tc.wantInterval, stream.pump.interval) + }) + } +} + +func TestAnthropicCompatStreaming_DisconnectDrainEndsWhenIntervalDisabled(t *testing.T) { + original := anthropicCompatDisconnectDrainInterval + anthropicCompatDisconnectDrainInterval = 200 * time.Millisecond + t.Cleanup(func() { anthropicCompatDisconnectDrainInterval = original }) + + for name, handle := range compatStreamHandlers() { + t.Run(name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + writer := &disconnectOnMarkerResponseWriter{ResponseRecorder: httptest.NewRecorder(), marker: "second"} + c, _ := gin.CreateTestContext(writer) + resp, pw := stallingAnthropicUpstream(t, anthropicStreamTruncatedBeforeMessageDelta()) + + got := runCompatStreamWithDeadline(t, handle, newCompatStreamIdleTestService(0), resp, c) + + require.NoError(t, got.err) + require.NotNil(t, got.result) + require.True(t, got.result.ClientDisconnect) + require.Equal(t, 20, got.result.Usage.InputTokens) + require.Zero(t, writer.writesAfterError) + _, werr := io.WriteString(pw, "event: ping\n") + require.Error(t, werr, "upstream body must be closed after the disconnect drain bound") + }) + } +} + +func TestForwardAsCompatStreaming_IdleTimeoutTriggersStreamTimeoutHandling(t *testing.T) { + for _, endpoint := range []string{"chat_completions", "responses"} { + for _, tc := range []struct { + name string + stream string + marker string + wantErr bool + wantOutput int + wantPenalty int + }{ + {name: "client_online", stream: anthropicStreamTruncatedBeforeMessageDelta(), marker: "\x00never-written", wantErr: true, wantOutput: 1, wantPenalty: 1}, + {name: "client_disconnected", stream: anthropicStreamTruncatedBeforeMessageDelta(), marker: "second", wantOutput: 1}, + {name: "after_message_stop", stream: anthropicStreamWithMidStreamText(), marker: "\x00never-written", wantOutput: 42}, + {name: "after_stop_reason", stream: anthropicStreamWithoutMessageStop(), marker: "\x00never-written", wantOutput: 42}, + } { + t.Run(endpoint+"/"+tc.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + writer := &disconnectOnMarkerResponseWriter{ResponseRecorder: httptest.NewRecorder(), marker: tc.marker} + c, _ := gin.CreateTestContext(writer) + body := `{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi"}],"stream":true}` + path := "/v1/chat/completions" + if endpoint == "responses" { + body = `{"model":"claude-sonnet-4-6","input":"hi","stream":true}` + path = "/v1/responses" + } + c.Request = httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)) + resp, _ := stallingAnthropicUpstream(t, tc.stream) + resp.StatusCode = http.StatusOK + repo := &errorPolicyRepoStub{} + settings := NewSettingService(&settingRepoStub{values: map[string]string{ + SettingKeyStreamTimeoutSettings: `{"enabled":true,"action":"temp_unsched","temp_unsched_minutes":1,"threshold_count":1,"threshold_window_minutes":1}`, + }}, nil) + cfg := rawChatCompletionsTestConfig() + cfg.Gateway.StreamDataIntervalTimeout = 1 + svc := &GatewayService{ + cfg: cfg, + httpUpstream: &anthropicHTTPUpstreamRecorder{resp: resp}, + rateLimitService: &RateLimitService{accountRepo: repo, settingService: settings}, + } + account := &Account{ID: 1, Platform: PlatformAnthropic, Type: AccountTypeAPIKey, Credentials: map[string]any{"api_key": "sk-test"}} + + var result *ForwardResult + var err error + if endpoint == "responses" { + result, err = svc.ForwardAsResponses(context.Background(), c, account, []byte(body), nil) + } else { + result, err = svc.ForwardAsChatCompletions(context.Background(), c, account, []byte(body), nil) + } + + if tc.wantErr { + require.ErrorIs(t, err, errAnthropicNativeStreamIdle) + } else { + require.NoError(t, err) + } + require.NotNil(t, result) + require.Equal(t, tc.wantOutput, result.Usage.OutputTokens) + require.Equal(t, tc.wantPenalty, repo.tempCalls) + }) + } + } +} + +func TestHandleCCStreamingFromAnthropic_ClientDisconnectDrainsUpstreamUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + writer := &disconnectOnMarkerResponseWriter{ResponseRecorder: httptest.NewRecorder(), marker: "second"} + c, _ := gin.CreateTestContext(writer) + + result, err := (&GatewayService{}).handleCCStreamingFromAnthropic( + &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader(anthropicStreamWithMidStreamText()))}, + c, "claude-sonnet-4-6", "claude-sonnet-4-6", nil, time.Now(), + ) + require.NoError(t, err) + require.NotNil(t, result) + require.True(t, writer.failed) + require.True(t, result.ClientDisconnect) + require.Zero(t, writer.writesAfterError) + require.Zero(t, writer.flushesAfterError) + + body := writer.Body.String() + require.Contains(t, body, `"content":"first"`) + require.NotContains(t, body, "second") + require.NotContains(t, body, "[DONE]") + + require.Equal(t, 20, result.Usage.InputTokens) + require.Equal(t, 42, result.Usage.OutputTokens) + require.Equal(t, 11, result.Usage.CacheReadInputTokens) + require.Equal(t, 4, result.Usage.CacheCreationInputTokens) + require.NotNil(t, result.FirstTokenMs) +} + +func TestHandleResponsesStreamingResponse_ClientDisconnectDrainsUpstreamUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + writer := &disconnectOnMarkerResponseWriter{ResponseRecorder: httptest.NewRecorder(), marker: "second"} + c, _ := gin.CreateTestContext(writer) + + result, err := (&GatewayService{}).handleResponsesStreamingResponse( + &http.Response{Header: http.Header{}, Body: io.NopCloser(strings.NewReader(anthropicStreamWithMidStreamText()))}, + c, "claude-sonnet-4-6", "claude-sonnet-4-6", nil, time.Now(), apicompat.ResponsesClientToolMapping{}, + ) + require.NoError(t, err) + require.NotNil(t, result) + require.True(t, writer.failed) + require.True(t, result.ClientDisconnect) + require.Zero(t, writer.writesAfterError) + require.Zero(t, writer.flushesAfterError) + + body := writer.Body.String() + require.Contains(t, body, `"delta":"first"`) + require.NotContains(t, body, "second") + require.NotContains(t, body, "response.completed") + + require.Equal(t, 20, result.Usage.InputTokens) + require.Equal(t, 42, result.Usage.OutputTokens) + require.Equal(t, 11, result.Usage.CacheReadInputTokens) + require.Equal(t, 4, result.Usage.CacheCreationInputTokens) + require.NotNil(t, result.FirstTokenMs) +} diff --git a/backend/internal/service/gateway_forward_compat_stream.go b/backend/internal/service/gateway_forward_compat_stream.go new file mode 100644 index 000000000000..7be3e85c35f2 --- /dev/null +++ b/backend/internal/service/gateway_forward_compat_stream.go @@ -0,0 +1,291 @@ +package service + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "github.com/tidwall/gjson" + "go.uber.org/zap" +) + +// anthropicCompatDisconnectDrainInterval 是 gateway.stream_data_interval_timeout +// 未启用时，客户端断开后继续读取上游的读间隔上限，与该配置项的默认值一致。 +var anthropicCompatDisconnectDrainInterval = 180 * time.Second + +// anthropicCompatStream 是 Anthropic 上游经 chat/responses 兼容入口流式转发时的 +// 上游读取状态。上游请求 ctx 与客户端解绑，客户端断开后调用方停止向下游写入、 +// 继续读取上游以合并最终 usage。逐行读间隔受 gateway.stream_data_interval_timeout +// 限制；该配置未启用时，客户端断开后以 anthropicCompatDisconnectDrainInterval 限制。 +type anthropicCompatStream struct { + pump *anthropicNativeLinePump + body io.Closer + maxLineSize int + logPrefix string + requestID string + // clientGone 是请求 ctx 的 Done，客户端断开后为 nil。 + clientGone <-chan struct{} + + clientDisconnected bool + sawMessageStart bool + // sawTerminal 表示上游已发出 message_stop、带 stop_reason 的 message_delta 或 [DONE]。 + sawTerminal bool + // upstreamErr* 记录上游流内的 error 事件；收到后调用方结束读取。 + upstreamErrType string + upstreamErrMessage string + upstreamErrPayload string + // readErr 是上游流结束的原因（io.EOF、读错误或 errAnthropicNativeStreamIdle）。 + readErr error +} + +// anthropicCompatStreamFailure 描述客户端在线时上游流未正常结束的原因： +// err 返回给调用方，errType / message 写入下游的终止错误事件。 +type anthropicCompatStreamFailure struct { + err error + errType string + message string +} + +// newAnthropicCompatStream 启动上游行泵；调用方 defer stop()。 +func newAnthropicCompatStream(cfg *config.Config, c *gin.Context, resp *http.Response, logPrefix, requestID string) *anthropicCompatStream { + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if cfg != nil && cfg.Gateway.MaxLineSize > 0 { + maxLineSize = cfg.Gateway.MaxLineSize + } + scanner.Buffer(make([]byte, 0, 64*1024), maxLineSize) + stream := &anthropicCompatStream{ + pump: newAnthropicNativeLinePump(scanner, anthropicNativeStreamInterval(cfg)), + body: resp.Body, + maxLineSize: maxLineSize, + logPrefix: logPrefix, + requestID: requestID, + } + if c.Request != nil { + stream.clientGone = c.Request.Context().Done() + } + return stream +} + +func (st *anthropicCompatStream) stop() { + st.pump.stop() +} + +// markClientDisconnected 标记客户端已断开，此后调用方不再向下游写入；读间隔上限 +// 未启用时以 anthropicCompatDisconnectDrainInterval 启用。 +func (st *anthropicCompatStream) markClientDisconnected() { + if st.clientDisconnected { + return + } + st.clientDisconnected = true + st.clientGone = nil + st.pump.enableInterval(anthropicCompatDisconnectDrainInterval) + logger.L().Info(st.logPrefix+" stream: client disconnected, continuing to drain upstream for billing", + zap.String("request_id", st.requestID), + ) +} + +// nextLine 返回上游下一行；ok=false 表示上游流已结束，原因记入 readErr。 +// 读间隔超时时关闭 body，解除泵 goroutine 阻塞中的读并释放上游连接。 +// 上游流结束时请求 ctx 已取消的，按客户端已断开收尾。 +func (st *anthropicCompatStream) nextLine() (string, bool) { + for { + line, err := st.pump.nextOrClientGone(st.clientGone) + if err == nil { + if payload, ok := extractOpenAISSEDataLine(line); ok && strings.TrimSpace(payload) == "[DONE]" { + st.sawTerminal = true + } + return line, true + } + if errors.Is(err, errAnthropicNativeClientGone) { + st.markClientDisconnected() + continue + } + if st.clientGone != nil { + select { + case <-st.clientGone: + st.clientDisconnected = true + st.clientGone = nil + default: + } + } + st.readErr = err + switch { + case errors.Is(err, errAnthropicNativeStreamIdle): + _ = st.body.Close() + logger.L().Warn(st.logPrefix+" stream: data interval timeout", + zap.String("request_id", st.requestID), + zap.Duration("interval", st.pump.interval), + zap.Bool("client_disconnected", st.clientDisconnected), + ) + case errors.Is(err, io.EOF), errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + // 正常收流或取消：不记日志。 + default: + logger.L().Warn(st.logPrefix+" stream: read error", + zap.Error(err), + zap.String("request_id", st.requestID), + ) + } + return "", false + } +} + +// observe 记录决定收尾方式的上游事件。收到上游 error 事件时关闭上游 body， +// 调用方不再处理该事件并结束读取。 +func (st *anthropicCompatStream) observe(event *apicompat.AnthropicStreamEvent, payload string) { + switch event.Type { + case "message_start": + st.sawMessageStart = true + case "message_stop": + st.sawTerminal = true + case "message_delta": + if event.Delta != nil && event.Delta.StopReason != "" { + st.sawTerminal = true + } + case "error": + st.upstreamErrPayload = payload + st.upstreamErrType = firstNonEmpty(gjson.Get(payload, "error.type").String(), "upstream_error") + st.upstreamErrMessage = firstNonEmpty(sanitizeUpstreamErrorMessage(gjson.Get(payload, "error.message").String()), "upstream stream error") + _ = st.body.Close() + logger.L().Warn(st.logPrefix+" stream: upstream error event", + zap.String("request_id", st.requestID), + zap.String("error_type", st.upstreamErrType), + zap.String("error_message", st.upstreamErrMessage), + zap.Bool("client_disconnected", st.clientDisconnected), + ) + } +} + +// failure 返回客户端在线、上游未发出结束事件即中断时的失败原因：上游流内 +// error 事件、读间隔超时、读错误，或 message_start 之后正常关闭。此时调用方不补发 +// 完成事件，而是写出协议内的终止错误。上游已发出结束事件、客户端已断开，或上游 +// 未发出 message_start 即正常关闭时返回 nil，调用方按已累计的 usage 正常收尾。 +// 读间隔超时时 err 为 errAnthropicNativeStreamIdle。 +func (st *anthropicCompatStream) failure() *anthropicCompatStreamFailure { + if st.clientDisconnected || st.sawTerminal { + return nil + } + switch { + case st.upstreamErrType != "": + return &anthropicCompatStreamFailure{ + err: &sseStreamErrorEventError{RawData: st.upstreamErrPayload}, + errType: st.upstreamErrType, + message: st.upstreamErrMessage, + } + case errors.Is(st.readErr, errAnthropicNativeStreamIdle): + return &anthropicCompatStreamFailure{ + err: errAnthropicNativeStreamIdle, + errType: "stream_timeout", + message: fmt.Sprintf("upstream stream idle for %s", st.pump.interval), + } + case errors.Is(st.readErr, bufio.ErrTooLong): + return &anthropicCompatStreamFailure{ + err: fmt.Errorf("stream usage incomplete: %w", st.readErr), + errType: "response_too_large", + message: fmt.Sprintf("upstream SSE line exceeded %d bytes", st.maxLineSize), + } + case st.readErr != nil && !errors.Is(st.readErr, io.EOF): + return &anthropicCompatStreamFailure{ + err: fmt.Errorf("stream usage incomplete: %w", st.readErr), + errType: "stream_read_error", + message: "upstream stream disconnected: " + sanitizeStreamError(st.readErr), + } + case st.sawMessageStart: + return &anthropicCompatStreamFailure{ + err: errors.New("stream usage incomplete: missing terminal event"), + errType: "stream_incomplete", + message: "upstream stream ended without a terminal event", + } + default: + return nil + } +} + +// finishAnthropicCompatStreamFailure 按失败原因收尾，返回交给调用方的结果与错误。 +// 上游流内 error 事件与 /v1/messages 一致：不计费，返回 sseStreamErrorEventError +// 由调用方转为 UpstreamFailoverError；尚未向客户端写出内容时不写终止错误，由 +// handler 换号重试。其余失败写出终止错误，并带上已观测到的 usage；未观测到任何 +// usage 时结果为 nil。 +func finishAnthropicCompatStreamFailure(c *gin.Context, writerSizeBeforeStream int, failure *anthropicCompatStreamFailure, result *ForwardResult, writeFailure func()) (*ForwardResult, error) { + var sseErr *sseStreamErrorEventError + if errors.As(failure.err, &sseErr) { + if c.Writer.Size() != writerSizeBeforeStream { + writeFailure() + } + return nil, failure.err + } + writeFailure() + if !result.Usage.hasObservedTokens() { + return nil, failure.err + } + return result, failure.err +} + +// writeAnthropicCompatCCStreamFailure 向 Chat Completions 流写出终止错误帧， +// 并标记响应已提交，handler 不再追加兜底错误。 +func writeAnthropicCompatCCStreamFailure(c *gin.Context, failure *anthropicCompatStreamFailure) { + payload, _ := json.Marshal(map[string]any{ + "error": map[string]string{ + "type": failure.errType, + "message": failure.message, + }, + }) + commitAnthropicCompatStreamFailure(c, failure, "data: "+string(payload)+"\n\n") +} + +// writeAnthropicCompatResponsesStreamFailure 向 Responses 流写出 response.failed +// 终止事件，沿用流内的 response id；sequence_number 经 restorer 与已写出的事件 +// 保持连续。写出后标记响应已提交。 +func writeAnthropicCompatResponsesStreamFailure(c *gin.Context, state *apicompat.AnthropicEventToResponsesState, restorer *apicompat.ResponsesClientToolStreamRestorer, failure *anthropicCompatStreamFailure) { + responseID := state.ResponseID + if responseID == "" { + responseID = "resp_" + strings.ReplaceAll(uuid.NewString(), "-", "") + } + evt := apicompat.ResponsesStreamEvent{ + Type: "response.failed", + SequenceNumber: state.SequenceNumber, + Response: &apicompat.ResponsesResponse{ + ID: responseID, + Object: "response", + CreatedAt: state.Created, + Model: state.Model, + Status: "failed", + Output: []apicompat.ResponsesOutput{}, + Error: &apicompat.ResponsesError{Code: failure.errType, Message: failure.message}, + }, + } + state.SequenceNumber++ + payload, err := json.Marshal(evt) + if err != nil { + return + } + payloads, _, err := restorer.RestoreEvent(payload) + if err != nil { + return + } + var frames strings.Builder + for _, restored := range payloads { + fmt.Fprintf(&frames, "event: %s\ndata: %s\n\n", gjson.GetBytes(restored, "type").String(), restored) + } + commitAnthropicCompatStreamFailure(c, failure, frames.String()) +} + +func commitAnthropicCompatStreamFailure(c *gin.Context, failure *anthropicCompatStreamFailure, frame string) { + MarkOpsStreamError(c, failure.errType, failure.message, http.StatusBadGateway) + MarkResponseCommitted(c) + if _, err := io.WriteString(c.Writer, frame); err != nil { + return + } + c.Writer.Flush() +} diff --git a/backend/internal/service/gateway_forward_partial_usage_test.go b/backend/internal/service/gateway_forward_partial_usage_test.go index 617940ce9d13..90f6757e0119 100644 --- a/backend/internal/service/gateway_forward_partial_usage_test.go +++ b/backend/internal/service/gateway_forward_partial_usage_test.go @@ -261,7 +261,7 @@ func TestGatewayService_Forward_PreOutputSSEOverloadedErrorUsesSemantic529(t *te require.Empty(t, rec.Body.String(), "pre-output overload must remain eligible for account failover") } -func TestGatewayService_Forward_PostOutputSSEOverloadedErrorUsesSemantic529(t *testing.T) { +func TestGatewayService_Forward_PostOutputSSEOverloadedErrorKeepsExistingStatus(t *testing.T) { gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() @@ -274,7 +274,6 @@ func TestGatewayService_Forward_PostOutputSSEOverloadedErrorUsesSemantic529(t *t const errorJSON = `{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}` fixture := "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n" + - "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n" + "event: error\ndata: " + errorJSON + "\n\n" upstream := &anthropicHTTPUpstreamRecorder{resp: &http.Response{ StatusCode: http.StatusOK, @@ -297,13 +296,10 @@ func TestGatewayService_Forward_PostOutputSSEOverloadedErrorUsesSemantic529(t *t var failoverErr *UpstreamFailoverError require.ErrorAs(t, err, &failoverErr) - require.Equal(t, 529, failoverErr.StatusCode, "stream output must not change the upstream error classification") + require.Equal(t, http.StatusForbidden, failoverErr.StatusCode) require.JSONEq(t, errorJSON, string(failoverErr.ResponseBody)) require.Zero(t, repo.tempCalls) - require.Zero(t, repo.overloadCalls, "a post-output overload must not add the pre-output failover cooldown") - require.Empty(t, repo.modelRateLimitCalls) require.Contains(t, rec.Body.String(), "message_start") - require.Contains(t, rec.Body.String(), "content_block_start") } func TestGatewayService_AnthropicAPIKeyPassthrough_ForwardStreamMissingTerminalPreservesPartialUsage(t *testing.T) { diff --git a/backend/internal/service/gateway_scheduling.go b/backend/internal/service/gateway_scheduling.go index 5d497bd118f7..516d476caa04 100644 --- a/backend/internal/service/gateway_scheduling.go +++ b/backend/internal/service/gateway_scheduling.go @@ -1122,10 +1122,22 @@ func (s *GatewayService) listSchedulableAccounts(ctx context.Context, groupID *i return accounts, useMixed, nil } -// IsSingleAntigravityAccountGroup 检查指定分组是否只有一个 antigravity 平台的可调度账号。 -// 用于 Handler 层在首次请求时提前设置 SingleAccountRetry context， -// 避免单账号分组收到 503 时错误地设置模型限流标记导致后续请求连续快速失败。 -func (s *GatewayService) IsSingleAntigravityAccountGroup(ctx context.Context, groupID *int64) bool { +// WithSingleAntigravityAccountGroupProbe 为请求挂载「分组是否只有一个 antigravity 可调度账号」的延迟判定。 +// 只有 antigravity 账号遇到 503 或模型限流预检查时才会查询；其他平台的请求不读取该判定。 +func (s *GatewayService) WithSingleAntigravityAccountGroupProbe(ctx context.Context, groupID *int64) context.Context { + var probeGroupID *int64 + if groupID != nil { + id := *groupID + probeGroupID = &id + } + return WithSingleAccountRetryProbe(ctx, func(probeCtx context.Context) bool { + return s.isSingleAntigravityAccountGroup(probeCtx, probeGroupID) + }) +} + +// isSingleAntigravityAccountGroup 检查指定分组是否只有一个 antigravity 平台的可调度账号。 +// 单账号分组收到 503 时原地重试而不设置模型限流，避免后续请求连续快速失败。 +func (s *GatewayService) isSingleAntigravityAccountGroup(ctx context.Context, groupID *int64) bool { accounts, _, err := s.listSchedulableAccounts(ctx, groupID, PlatformAntigravity, true) if err != nil { return false @@ -1309,9 +1321,9 @@ func (s *GatewayService) withWindowCostPrefetch(ctx context.Context, accounts [] } // isAccountSchedulableForQuota 检查账号是否在配额限制内 -// 适用于配置了 quota_limit 的 apikey 和 bedrock 类型账号 +// 适用于配置了 quota_limit 的 apikey、bedrock 和 service_account 类型账号 func (s *GatewayService) isAccountSchedulableForQuota(account *Account) bool { - if !account.IsAPIKeyOrBedrock() { + if !account.SupportsAccountQuota() { return true } return !account.IsQuotaExceeded() diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 6384920549f9..057c422c76f8 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -789,9 +789,6 @@ func (e *UpstreamFailoverError) ShouldReportAccountScheduleFailure() bool { // 提取 RawData 并构造 UpstreamFailoverError.ResponseBody。 type sseStreamErrorEventError struct { RawData string - // MessageStarted excludes keepalive pings/comments and HTTP headers. Once - // a message event reaches the client, another attempt cannot replace it. - MessageStarted bool } func (e *sseStreamErrorEventError) Error() string { return "have error in stream" } diff --git a/backend/internal/service/gateway_upstream_response.go b/backend/internal/service/gateway_upstream_response.go index 2ab9f881512a..254417116d44 100644 --- a/backend/internal/service/gateway_upstream_response.go +++ b/backend/internal/service/gateway_upstream_response.go @@ -666,11 +666,6 @@ type streamingResult struct { clientDisconnect bool // 客户端是否在流式传输过程中断开 } -// Only message_start is held back, so a signature rejection before content can -// be repaired without exposing two message IDs to the client. Larger preludes -// are committed immediately rather than retaining an unbounded upstream frame. -const claudeMessageStartPreludeLimit = 64 * 1024 - // hasObservedTokens 报告流式过程中是否已观测到任何上游计量的 token。 func (u *ClaudeUsage) hasObservedTokens() bool { if u == nil { @@ -713,7 +708,7 @@ func partialStreamUsageResult(c *gin.Context, resp *http.Response, streamResult } } -func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel string, mimicClaudeCode bool, allowPreludeRecovery ...bool) (*streamingResult, error) { +func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel string, mimicClaudeCode bool) (*streamingResult, error) { observer := upstreamResponseModelObserverFromContext(c) if observer == nil { observer = beginUpstreamResponseModelObservation(c) @@ -858,42 +853,6 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http needModelReplace := originalModel != mappedModel clientDisconnected := false // 客户端断开标志，断开后继续读取上游以获取完整usage - messageStarted := false // 本次尝试已提交非 ping 事件；HTTP headers/ping 不妨碍签名恢复。 - preludeRecoveryEnabled := len(allowPreludeRecovery) > 0 && allowPreludeRecovery[0] - // The second optional flag is used only by a sent signature repair. A - // failed repair belongs to the rejected request, not the account's health. - suppressAccountHealthPenalty := len(allowPreludeRecovery) > 1 && allowPreludeRecovery[1] - bufferPrelude := preludeRecoveryEnabled - var pendingPrelude []string - pendingPreludeBytes := 0 - writeBlocks := func(blocks []string, eventType string) { - for _, block := range blocks { - if clientDisconnected { - continue - } - restored := reverseToolNamesIfPresent(c, []byte(block)) - written, werr := fmt.Fprint(w, string(restored)) - if written > 0 && eventType != "ping" { - messageStarted = true - } - if werr != nil { - clientDisconnected = true - logger.LegacyPrintf("service.gateway", "Client disconnected during streaming, continuing to drain upstream for billing") - continue - } - flusher.Flush() - lastDataAt = time.Now() - resetKeepaliveTimer() - } - } - commitPrelude := func() { - if len(pendingPrelude) == 0 { - return - } - writeBlocks(pendingPrelude, "message_start") - pendingPrelude = nil - pendingPreludeBytes = 0 - } sawTerminalEvent := false useNoopDeltaKeepalive := c != nil && c.Request != nil && shouldUseClaudeCodeNoopDeltaKeepalive(c.GetHeader("User-Agent")) noopDeltaKeepaliveBlockIndex := -1 @@ -901,9 +860,9 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http pendingEventLines := make([]string, 0, 4) - processSSEEvent := func(lines []string) ([]string, string, *sseUsagePatch, string, error) { + processSSEEvent := func(lines []string) ([]string, string, *sseUsagePatch, error) { if len(lines) == 0 { - return nil, "", nil, "", nil + return nil, "", nil, nil } eventName := "" @@ -920,16 +879,11 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http } if eventName == "error" { - return nil, dataLine, nil, "error", &sseStreamErrorEventError{RawData: dataLine} + return nil, dataLine, nil, &sseStreamErrorEventError{RawData: dataLine} } if dataLine == "" { - // SSE comments are keepalives too; they must not commit a message prelude. - kind := eventName - if kind == "" { - kind = "ping" - } - return []string{strings.Join(lines, "\n") + "\n\n"}, "", nil, kind, nil + return []string{strings.Join(lines, "\n") + "\n\n"}, "", nil, nil } if dataLine == "[DONE]" { @@ -939,7 +893,7 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http block = "event: " + eventName + "\n" } block += "data: " + dataLine + "\n\n" - return []string{block}, dataLine, nil, "message_stop", nil + return []string{block}, dataLine, nil, nil } var event map[string]any @@ -950,29 +904,14 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http block = "event: " + eventName + "\n" } block += "data: " + dataLine + "\n\n" - kind := eventName - if kind == "" { - kind = "unknown" - } - return []string{block}, dataLine, nil, kind, nil + return []string{block}, dataLine, nil, nil } eventType, _ := event["type"].(string) - if eventType == "error" { - return nil, dataLine, nil, "error", &sseStreamErrorEventError{RawData: dataLine} - } observer.ObserveAnthropic([]byte(dataLine)) if eventName == "" { eventName = eventType } - eventKind := eventName - if eventType != "" && eventKind != eventType { - // A contradictory event name/payload is not a safe ping or prelude. - eventKind = "unknown" - } - if eventKind == "" { - eventKind = "unknown" - } eventChanged := false if useNoopDeltaKeepalive { @@ -1060,7 +999,7 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http block = "event: " + eventName + "\n" } block += "data: " + dataLine + "\n\n" - return []string{block}, dataLine, usagePatch, eventKind, nil + return []string{block}, dataLine, usagePatch, nil } newData, err := json.Marshal(event) @@ -1071,7 +1010,7 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http block = "event: " + eventName + "\n" } block += "data: " + dataLine + "\n\n" - return []string{block}, dataLine, usagePatch, eventKind, nil + return []string{block}, dataLine, usagePatch, nil } block := "" @@ -1079,29 +1018,20 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http block = "event: " + eventName + "\n" } block += "data: " + string(newData) + "\n\n" - return []string{block}, string(newData), usagePatch, eventKind, nil + return []string{block}, string(newData), usagePatch, nil } for { select { case ev, ok := <-events: if !ok { - hadPendingPrelude := len(pendingPrelude) > 0 - commitPrelude() // 上游完成，返回结果 if !sawTerminalEvent { - if hadPendingPrelude && !clientDisconnected { - // Committing the held start makes the response streaming; finish it - // explicitly rather than relying on a handler JSON fallback. - sendErrorEvent("upstream_error", "Upstream stream ended before completion") - MarkResponseCommitted(c) - } return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, fmt.Errorf("stream usage incomplete: missing terminal event") } return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, nil } if ev.err != nil { - commitPrelude() if sawTerminalEvent { return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, nil } @@ -1152,58 +1082,39 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http continue } - outputBlocks, data, usagePatch, eventType, err := processSSEEvent(pendingEventLines) + outputBlocks, data, usagePatch, err := processSSEEvent(pendingEventLines) pendingEventLines = pendingEventLines[:0] if err != nil { if clientDisconnected { return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil } - var sseErr *sseStreamErrorEventError - if errors.As(err, &sseErr) { - // Custom rectifier patterns share the same matcher as HTTP recovery. - requestRejected := ClassifyClaudeStreamError([]byte(data)).StatusCode == http.StatusBadRequest - signatureRejected := requestRejected && s.isThinkingBlockSignatureError([]byte(data)) - if requestRejected && !signatureRejected && bufferPrelude && s.settingService != nil { - signatureRejected = s.isSignatureErrorPattern(ctx, account, []byte(data)) - } - if bufferPrelude && signatureRejected && !messageStarted { - pendingPrelude = nil - pendingPreludeBytes = 0 - } else { - commitPrelude() - } - sseErr.MessageStarted = messageStarted - } - if usage.hasObservedTokens() { - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, err - } return nil, err } - if data != "" { - if firstTokenMs == nil && data != "[DONE]" && ((!preludeRecoveryEnabled && !suppressAccountHealthPenalty) || eventType != "ping") { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - if usagePatch != nil { - // Metering belongs to this attempt even while its prelude is held. - mergeSSEUsagePatch(usage, usagePatch) - } - } - if bufferPrelude && !messageStarted && len(pendingPrelude) == 0 && eventType == "message_start" { - pendingPrelude = append(pendingPrelude, outputBlocks...) - for _, block := range outputBlocks { - pendingPreludeBytes += len(block) - } - if pendingPreludeBytes > claudeMessageStartPreludeLimit { - commitPrelude() - bufferPrelude = false + for _, block := range outputBlocks { + if !clientDisconnected { + restored := reverseToolNamesIfPresent(c, []byte(block)) + if _, werr := fmt.Fprint(w, string(restored)); werr != nil { + clientDisconnected = true + logger.LegacyPrintf("service.gateway", "Client disconnected during streaming, continuing to drain upstream for billing") + // 不 break：客户端断开后仍需继续合并本事件及后续事件的 usage， + // 否则会漏计当前事件携带的 usage 导致少计费。后续写入由 + // clientDisconnected 守卫跳过。 + } else { + flusher.Flush() + lastDataAt = time.Now() + resetKeepaliveTimer() + } } - } else { - if eventType != "ping" { - commitPrelude() + if data != "" { + if firstTokenMs == nil && data != "[DONE]" { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + if usagePatch != nil { + mergeSSEUsagePatch(usage, usagePatch) + } } - writeBlocks(outputBlocks, eventType) } continue } @@ -1215,13 +1126,12 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http if time.Since(lastRead) < streamInterval { continue } - commitPrelude() if clientDisconnected { return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, fmt.Errorf("stream usage incomplete after timeout") } logger.LegacyPrintf("service.gateway", "Stream data interval timeout: account=%d model=%s interval=%s", account.ID, originalModel, streamInterval) // 处理流超时，可能标记账户为临时不可调度或错误状态 - if s.rateLimitService != nil && !suppressAccountHealthPenalty { + if s.rateLimitService != nil { s.rateLimitService.HandleStreamTimeout(ctx, account, originalModel) } sendErrorEvent("stream_timeout", fmt.Sprintf("upstream stream idle for %s", streamInterval)) diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index f05c11106038..66c3370e43b2 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -137,7 +137,7 @@ func (p *postUsageBillingParams) shouldUpdateRateLimits() bool { } func (p *postUsageBillingParams) shouldUpdateAccountQuota() bool { - return p.Cost.TotalCost > 0 && p.Account.IsAPIKeyOrBedrock() && p.Account.HasAnyQuotaLimit() + return p.Cost.TotalCost > 0 && p.Account.SupportsAccountQuota() && p.Account.HasAnyQuotaLimit() } // postUsageBilling is the legacy fallback billing path used when the unified @@ -348,6 +348,8 @@ func buildUsageBillingCommand(requestID string, usageLog *UsageLog, p *postUsage } if p.shouldUpdateAccountQuota() { cmd.AccountQuotaCost = p.Cost.TotalCost * p.AccountRateMultiplier + cmd.AccountQuotaHasDailyLimit = p.Account.GetQuotaDailyLimit() > 0 + cmd.AccountQuotaHasWeeklyLimit = p.Account.GetQuotaWeeklyLimit() > 0 } cmd.Normalize() @@ -550,11 +552,12 @@ func notifyAccountQuota(p *postUsageBillingParams, deps *billingDeps, result *Us slog.Error("panic in notifyAccountQuota", "recover", r) } }() - if p.Cost.TotalCost <= 0 || p.Account == nil || !p.Account.IsAPIKeyOrBedrock() || deps.balanceNotifyService == nil { + if p.Cost.TotalCost <= 0 || p.Account == nil || !p.Account.SupportsAccountQuota() || !p.Account.HasAnyQuotaLimit() || deps.balanceNotifyService == nil { slog.Debug("notifyAccountQuota: skipped", "total_cost", p.Cost.TotalCost, "account_nil", p.Account == nil, - "is_apikey_or_bedrock", p.Account != nil && p.Account.IsAPIKeyOrBedrock(), + "supports_account_quota", p.Account != nil && p.Account.SupportsAccountQuota(), + "has_quota_limit", p.Account != nil && p.Account.HasAnyQuotaLimit(), "service_nil", deps.balanceNotifyService == nil, ) return @@ -563,6 +566,10 @@ func notifyAccountQuota(p *postUsageBillingParams, deps *billingDeps, result *Us var quotaState *AccountQuotaState if result != nil { quotaState = result.QuotaState + // 本应累加却无递增后状态（账号已删除，计数被跳过）：没有新用量，不做阈值判定。 + if quotaState == nil && p.shouldUpdateAccountQuota() { + return + } } slog.Debug("notifyAccountQuota: calling CheckAccountQuotaAfterIncrement", "account_id", p.Account.ID, diff --git a/backend/internal/service/gateway_usage_billing_account_quota_test.go b/backend/internal/service/gateway_usage_billing_account_quota_test.go new file mode 100644 index 000000000000..8d45bcc54c5e --- /dev/null +++ b/backend/internal/service/gateway_usage_billing_account_quota_test.go @@ -0,0 +1,92 @@ +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBuildUsageBillingCommand_AccountQuotaPeriodHints(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + accountType string + extra map[string]any + wantQuota bool + wantDaily bool + wantWeekly bool + }{ + {name: "no_limits", extra: map[string]any{}}, + {name: "total_only", extra: map[string]any{"quota_limit": 10.0}, wantQuota: true}, + {name: "daily_only", extra: map[string]any{"quota_daily_limit": 5.0}, wantQuota: true, wantDaily: true}, + {name: "weekly_only", extra: map[string]any{"quota_weekly_limit": 5.0}, wantQuota: true, wantWeekly: true}, + { + name: "all_dimensions", + extra: map[string]any{"quota_limit": 10.0, "quota_daily_limit": 5.0, "quota_weekly_limit": 7.0}, + wantQuota: true, wantDaily: true, wantWeekly: true, + }, + { + name: "non_positive_period_limits", + extra: map[string]any{"quota_limit": 10.0, "quota_daily_limit": 0.0, "quota_weekly_limit": -1.0}, + wantQuota: true, + }, + { + name: "service_account_total_only", + accountType: AccountTypeServiceAccount, + extra: map[string]any{"quota_limit": 10.0}, + wantQuota: true, + }, + { + name: "service_account_daily", + accountType: AccountTypeServiceAccount, + extra: map[string]any{"quota_limit": 10.0, "quota_daily_limit": 5.0}, + wantQuota: true, wantDaily: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + accountType := tt.accountType + if accountType == "" { + accountType = AccountTypeAPIKey + } + p := &postUsageBillingParams{ + Cost: &CostBreakdown{TotalCost: 1, ActualCost: 1}, + User: &User{ID: 1}, + APIKey: &APIKey{ID: 2}, + Account: &Account{ID: 3, Type: accountType, Extra: tt.extra}, + AccountRateMultiplier: 1, + } + + cmd := buildUsageBillingCommand("req-1", nil, p) + require.NotNil(t, cmd) + require.Equal(t, tt.wantQuota, cmd.AccountQuotaCost > 0) + require.Equal(t, tt.wantDaily, cmd.AccountQuotaHasDailyLimit) + require.Equal(t, tt.wantWeekly, cmd.AccountQuotaHasWeeklyLimit) + }) + } +} + +func TestUsageBillingCommandFingerprintIgnoresAccountQuotaPeriodHints(t *testing.T) { + t.Parallel() + + base := UsageBillingCommand{ + RequestID: "req-1", + UserID: 1, + APIKeyID: 2, + AccountID: 3, + AccountType: AccountTypeAPIKey, + BalanceCost: 1, + AccountQuotaCost: 1, + } + hinted := base + hinted.AccountQuotaHasDailyLimit = true + hinted.AccountQuotaHasWeeklyLimit = true + + base.Normalize() + hinted.Normalize() + require.NotEmpty(t, base.RequestFingerprint) + require.Equal(t, base.RequestFingerprint, hinted.RequestFingerprint) +} diff --git a/backend/internal/service/gemini_chat_completions_compat_service.go b/backend/internal/service/gemini_chat_completions_compat_service.go index 33304d5c856f..d9ace5909269 100644 --- a/backend/internal/service/gemini_chat_completions_compat_service.go +++ b/backend/internal/service/gemini_chat_completions_compat_service.go @@ -100,6 +100,9 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( if err != nil { return nil, s.writeChatCompletionsError(c, http.StatusBadRequest, "invalid_request_error", err.Error()) } + if geminiImageGenerationDenied(c, geminiReq, originalModel, req.Model, mappedModel) { + return nil, s.writeChatCompletionsError(c, http.StatusBadRequest, "invalid_request_error", ImageGenerationPermissionMessage()) + } geminiReq = ensureGeminiFunctionCallThoughtSignatures(geminiReq) geminiReq = applyGeminiCompatThinkingConfig(geminiReq, claudeBody, mappedModel, clientStream) geminiReq = s.systemPrompts.ApplyGemini(c, account, geminiReq) @@ -251,6 +254,8 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( var usage *ClaudeUsage var firstTokenMs *int + clientDisconnect := false + streamFinished := false if clientStream { streamRes, err := s.handleChatCompletionsStreamingResponseFromGemini(c, resp, startTime, originalModel, account.Type == AccountTypeOAuth, includeUsage) if err != nil { @@ -258,6 +263,8 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( } usage = streamRes.usage firstTokenMs = streamRes.firstTokenMs + clientDisconnect = streamRes.clientDisconnect + streamFinished = streamRes.finished } else if useUpstreamStream { collected, usageObj, err := collectGeminiSSE(resp.Body, account.Type == AccountTypeOAuth) if err != nil { @@ -288,6 +295,10 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( if isImageGenerationModel(originalModel) { imageCount = 1 } + // 客户端在收到结束分片前断开时只计已送达的图片：Chat 流不向客户端转发图片。 + if clientDisconnect && !streamFinished { + imageCount = 0 + } return &ForwardResult{ RequestID: requestID, @@ -302,7 +313,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( ImageCount: imageCount, ImageSize: imageSize, ImageInputSize: imageInputSize, - ClientDisconnect: false, + ClientDisconnect: clientDisconnect, }, nil } @@ -533,7 +544,11 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom var usage ClaudeUsage var firstTokenMs *int firstChunk := true - var protocolErr error + finishReason := "" + clientDisconnect := false + streamResult := func() *geminiStreamResult { + return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnect, finished: finishReason != ""} + } writeChatChunk := func(chunk apicompat.ChatCompletionsChunk) bool { sse, err := apicompat.ChatChunkToSSE(chunk) @@ -541,39 +556,20 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom return false } if _, err := io.WriteString(c.Writer, sse); err != nil { + clientDisconnect = true return true } return false } - writeResponseEvent := func(event *apicompat.ResponsesStreamEvent) bool { - if protocolErr != nil { - return false - } - chunks := apicompat.ResponsesEventToChatChunks(event, ccState) - if ccState.ProtocolError != "" { - protocolErr = errors.New(ccState.ProtocolError) - MarkResponseCommitted(c) - _, _ = io.WriteString(c.Writer, buildChatStreamErrorSSE("upstream_protocol_error", ccState.ProtocolError)+"data: [DONE]\n\n") - flusher.Flush() - return false - } - for _, chunk := range chunks { - if writeChatChunk(chunk) { - return true - } - } - return false - } - emitAnthropicEvent := func(evt *apicompat.AnthropicStreamEvent) bool { - if protocolErr != nil { - return false // Drain usage, but do not convert or emit after the error. - } responsesEvents := apicompat.AnthropicEventToResponsesEvents(evt, anthState) for _, resEvt := range responsesEvents { - if writeResponseEvent(&resEvt) { - return true + chunks := apicompat.ResponsesEventToChatChunks(&resEvt, ccState) + for _, chunk := range chunks { + if disconnected := writeChatChunk(chunk); disconnected { + return true + } } } flusher.Flush() @@ -593,10 +589,9 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom Usage: apicompat.AnthropicUsage{}, }, }) { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } - finishReason := "" sawToolUse := false nextBlockIndex := 0 openBlockIndex := -1 @@ -665,12 +660,12 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom } if openToolIndex >= 0 { if closeOpenTool() { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } } if openBlockType != "thinking" { if closeOpenBlock() { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } idx := nextBlockIndex nextBlockIndex++ @@ -681,21 +676,21 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom Index: &idx, ContentBlock: &apicompat.AnthropicContentBlock{Type: "thinking"}, }) { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } } if emitAnthropicEvent(&apicompat.AnthropicStreamEvent{ Type: "content_block_delta", Delta: &apicompat.AnthropicDelta{Type: "thinking_delta", Thinking: delta}, }) { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } continue } if text, ok := part["text"].(string); ok && text != "" { if openToolIndex >= 0 { if closeOpenTool() { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } } delta, newSeen := computeGeminiTextDelta(seenText, text) @@ -705,7 +700,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom } if openBlockType != "text" { if closeOpenBlock() { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } idx := nextBlockIndex nextBlockIndex++ @@ -719,7 +714,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom Text: "", }, }) { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } } if emitAnthropicEvent(&apicompat.AnthropicStreamEvent{ @@ -729,7 +724,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom Text: delta, }, }) { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } continue } @@ -740,11 +735,11 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom name = "tool" } if closeOpenBlock() { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } if openToolIndex >= 0 && openToolName != name { if closeOpenTool() { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } } if openToolIndex < 0 { @@ -763,7 +758,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom Input: json.RawMessage(`{}`), }, }) { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } } @@ -789,7 +784,7 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom PartialJSON: delta, }, }) { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } } } @@ -803,18 +798,20 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom break } if err != nil { - if protocolErr != nil { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, protocolErr + if geminiStreamAbortedWithBillableOutput(c, err, &usage) { + logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] Client disconnected during chat completions stream, billing collected usage: %v", err) + clientDisconnect = true + return streamResult(), nil } return nil, fmt.Errorf("stream read error: %w", err) } } if closeOpenBlock() { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } if closeOpenTool() { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } stopReason := mapGeminiFinishReasonToClaudeStopReason(finishReason) @@ -835,30 +832,30 @@ func (s *GeminiMessagesCompatService) handleChatCompletionsStreamingResponseFrom CacheReadInputTokens: usage.CacheReadInputTokens, }, }) { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } if emitAnthropicEvent(&apicompat.AnthropicStreamEvent{Type: "message_stop"}) { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } for _, resEvt := range apicompat.FinalizeAnthropicResponsesStream(anthState) { - if writeResponseEvent(&resEvt) { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, protocolErr + chunks := apicompat.ResponsesEventToChatChunks(&resEvt, ccState) + for _, chunk := range chunks { + if disconnected := writeChatChunk(chunk); disconnected { + return streamResult(), nil + } } } - if protocolErr != nil { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, protocolErr - } for _, chunk := range apicompat.FinalizeResponsesChatStream(ccState) { if disconnected := writeChatChunk(chunk); disconnected { - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } } _, _ = io.WriteString(c.Writer, "data: [DONE]\n\n") flusher.Flush() - return &geminiStreamResult{usage: &usage, firstTokenMs: firstTokenMs}, nil + return streamResult(), nil } func (s *GeminiMessagesCompatService) writeGeminiChatCompletionsMappedError( diff --git a/backend/internal/service/gemini_image_generation_gate.go b/backend/internal/service/gemini_image_generation_gate.go new file mode 100644 index 000000000000..89230836856f --- /dev/null +++ b/backend/internal/service/gemini_image_generation_gate.go @@ -0,0 +1,94 @@ +package service + +import ( + "slices" + "strings" + + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +// geminiImageModalityEnum 是 proto JSON 中 Modality.IMAGE 的整数枚举值。 +const geminiImageModalityEnum = 2 + +// GeminiImageGenerationDenied 判断 Gemini 系请求是否因分组关闭生图而必须拒绝。 +// 门禁只作用于后台展示该开关的 Gemini 系分组（gemini / antigravity / composite）； +// 无分组的 Key 与其他平台分组不受影响。models 中任一名字属于生图模型， +// 或请求体要求输出 IMAGE 模态，即视为生图请求。 +func GeminiImageGenerationDenied(group *Group, geminiBody []byte, models ...string) bool { + if group == nil || group.AllowImageGeneration { + return false + } + switch group.Platform { + case PlatformGemini, PlatformAntigravity, PlatformComposite: + default: + return false + } + return isGeminiImageGenerationIntent(geminiBody, models...) +} + +// geminiImageGenerationDenied 是转发层在账号模型映射之后的门禁：models 传入 +// 转发时可见的模型名（入参模型名与账号映射后的上游模型名）。命中时把请求标记为 +// 本地功能门禁，调用方负责按各自协议写出 400。 +func geminiImageGenerationDenied(c *gin.Context, geminiBody []byte, models ...string) bool { + if c == nil || !GeminiImageGenerationDenied(apiKeyGroup(getAPIKeyFromContext(c)), geminiBody, models...) { + return false + } + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) + return true +} + +func isGeminiImageGenerationIntent(geminiBody []byte, models ...string) bool { + return slices.ContainsFunc(models, isGeminiImageOutputModel) || geminiRequestAsksForImageOutput(geminiBody) +} + +// isGeminiImageOutputModel 匹配以图片为输出的 Gemini 模型名，与图片接口共用 +// isGeminiCompatibleImageModel 判定。 +func isGeminiImageOutputModel(model string) bool { + return isGeminiCompatibleImageModel(strings.TrimPrefix(strings.ToLower(strings.TrimSpace(model)), "models/")) +} + +// geminiRequestAsksForImageOutput 判断 generationConfig.responseModalities 是否包含 IMAGE。 +// Gemini REST 同时接受 camelCase 与 snake_case 字段名、字符串与整数枚举； +// 重复出现的键逐个检查。 +func geminiRequestAsksForImageOutput(body []byte) bool { + if len(body) == 0 || !gjson.ValidBytes(body) { + return false + } + found := false + gjson.ParseBytes(body).ForEach(func(key, config gjson.Result) bool { + if key.Str == "generationConfig" || key.Str == "generation_config" { + found = geminiGenerationConfigAsksForImage(config) + } + return !found + }) + return found +} + +func geminiGenerationConfigAsksForImage(config gjson.Result) bool { + if !config.IsObject() { + return false + } + found := false + config.ForEach(func(key, modalities gjson.Result) bool { + if key.Str == "responseModalities" || key.Str == "response_modalities" { + found = geminiModalitiesContainImage(modalities) + } + return !found + }) + return found +} + +func geminiModalitiesContainImage(modalities gjson.Result) bool { + found := false + modalities.ForEach(func(_, modality gjson.Result) bool { + switch modality.Type { + case gjson.String: + found = strings.EqualFold(strings.TrimSpace(modality.Str), "IMAGE") + case gjson.Number: + found = modality.Num == geminiImageModalityEnum + } + return !found + }) + return found +} diff --git a/backend/internal/service/gemini_image_generation_gate_test.go b/backend/internal/service/gemini_image_generation_gate_test.go new file mode 100644 index 000000000000..c384563127fa --- /dev/null +++ b/backend/internal/service/gemini_image_generation_gate_test.go @@ -0,0 +1,263 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestIsGeminiImageGenerationIntent(t *testing.T) { + cases := []struct { + name string + body string + models []string + want bool + }{ + {name: "pro image", models: []string{"gemini-3-pro-image"}, want: true}, + {name: "flash image preview with prefix", models: []string{"models/gemini-3.1-flash-image-preview"}, want: true}, + {name: "image generation experimental", models: []string{"gemini-2.0-flash-exp-image-generation"}, want: true}, + {name: "imageless suffix", models: []string{"gemini-3-pro-imageless"}, want: false}, + {name: "image without gemini prefix", models: []string{"unknown-image", "my-image-chat"}, want: false}, + {name: "banana alias without resolved image model", models: []string{"nano-banana-2"}, want: false}, + {name: "any resolved name", models: []string{"my-alias", "gemini-2.5-flash-image"}, want: true}, + {name: "text model", models: []string{"gemini-3.7-flash", "gemini-3.1-pro-preview"}, want: false}, + {name: "empty", models: []string{"", " "}, want: false}, + {name: "camelCase modalities", body: `{"generationConfig":{"responseModalities":["TEXT","IMAGE"]}}`, models: []string{"gemini-3.7-flash"}, want: true}, + {name: "snake_case modalities", body: `{"generation_config":{"response_modalities":["image"]}}`, want: true}, + {name: "mixed case keys", body: `{"generationConfig":{"response_modalities":[" Image "]}}`, want: true}, + {name: "text only modalities", body: `{"generationConfig":{"responseModalities":["TEXT"]}}`, models: []string{"gemini-3.7-flash"}, want: false}, + {name: "modalities string not array", body: `{"generationConfig":{"responseModalities":"IMAGE"}}`, want: true}, + {name: "no generation config", body: `{"contents":[{"parts":[{"text":"IMAGE"}]}]}`, want: false}, + {name: "integer enum modalities", body: `{"generationConfig":{"responseModalities":[1,2]}}`, want: true}, + {name: "integer enum text only", body: `{"generationConfig":{"responseModalities":[1]}}`, want: false}, + {name: "duplicate generationConfig keys", body: `{"generationConfig":{"responseModalities":["TEXT"]},"generationConfig":{"responseModalities":["IMAGE"]}}`, want: true}, + {name: "duplicate modalities keys", body: `{"generationConfig":{"responseModalities":["TEXT"],"responseModalities":["IMAGE"]}}`, want: true}, + {name: "invalid json", body: `{"generationConfig":{"responseModalities":["IMAGE"]`, want: false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, isGeminiImageGenerationIntent([]byte(tc.body), tc.models...)) + }) + } +} + +func TestGeminiImageGenerationDenied_GroupScope(t *testing.T) { + cases := []struct { + name string + group *Group + want bool + }{ + {name: "ungrouped key", group: nil, want: false}, + {name: "gemini disabled", group: &Group{Platform: PlatformGemini}, want: true}, + {name: "antigravity disabled", group: &Group{Platform: PlatformAntigravity}, want: true}, + {name: "composite disabled", group: &Group{Platform: PlatformComposite}, want: true}, + {name: "gemini enabled", group: &Group{Platform: PlatformGemini, AllowImageGeneration: true}, want: false}, + {name: "anthropic group without the switch", group: &Group{Platform: PlatformAnthropic}, want: false}, + {name: "openai group gated elsewhere", group: &Group{Platform: PlatformOpenAI}, want: false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, GeminiImageGenerationDenied(tc.group, nil, "gemini-3-pro-image")) + }) + } +} + +func withGeminiImageGateAPIKey(c *gin.Context, group *Group) { + c.Set("api_key", &APIKey{ID: 1, Group: group}) +} + +func imageDisabledGroup() *Group { + return &Group{ID: 10, Platform: PlatformGemini, AllowImageGeneration: false} +} + +func requireGeminiImagePermissionError(t *testing.T, rec *httptest.ResponseRecorder, wantGoogleStatus string) { + t.Helper() + require.Equal(t, http.StatusBadRequest, rec.Code) + var payload map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &payload)) + errObj, ok := payload["error"].(map[string]any) + require.True(t, ok, "error object: %s", rec.Body.String()) + require.Equal(t, ImageGenerationPermissionMessage(), errObj["message"]) + if wantGoogleStatus != "" { + require.Equal(t, wantGoogleStatus, errObj["status"]) + } else { + require.Equal(t, "invalid_request_error", errObj["type"]) + } +} + +func TestGeminiForwardNative_ImageGenerationDisabledGroupRejectsBeforeUpstream(t *testing.T) { + gin.SetMode(gin.TestMode) + textBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"draw a cat"}]}]}`) + imageModalityBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"draw a cat"}]}],"generationConfig":{"responseModalities":["TEXT","IMAGE"]}}`) + + cases := []struct { + name string + model string + mapping map[string]any + body []byte + }{ + {name: "image model", model: "gemini-3-pro-image", body: textBody}, + {name: "account mapping alias to image model", model: "my-draw", mapping: map[string]any{"my-draw": "gemini-3.1-flash-image"}, body: textBody}, + {name: "text model asking for image modality", model: "gemini-3.7-flash", body: imageModalityBody}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + svc, httpStub := newGeminiVertexAPIKeyTestService() + c, rec := newGeminiNativeTestContext(t) + withGeminiImageGateAPIKey(c, imageDisabledGroup()) + account := geminiVertexAPIKeyTestAccount("") + if tc.mapping != nil { + account.Credentials["model_mapping"] = tc.mapping + } + + result, err := svc.ForwardNative(context.Background(), c, account, tc.model, "generateContent", false, tc.body) + require.Error(t, err) + require.Nil(t, result) + require.Zero(t, httpStub.calls) + requireGeminiImagePermissionError(t, rec, "INVALID_ARGUMENT") + }) + } +} + +func TestGeminiForwardNative_ImageGenerationGateAllowsPermittedRequests(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"contents":[{"role":"user","parts":[{"text":"draw a cat"}]}]}`) + + cases := []struct { + name string + group *Group + model string + action string + }{ + {name: "group allows image generation", group: &Group{ID: 34, Platform: PlatformGemini, AllowImageGeneration: true}, model: "gemini-3-pro-image", action: "generateContent"}, + {name: "ungrouped key", group: nil, model: "gemini-3-pro-image", action: "generateContent"}, + {name: "text model in disabled group", group: imageDisabledGroup(), model: "gemini-3.7-flash", action: "generateContent"}, + {name: "countTokens on image model", group: imageDisabledGroup(), model: "gemini-3-pro-image", action: "countTokens"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + svc, httpStub := newGeminiVertexAPIKeyTestService() + c, rec := newGeminiNativeTestContext(t) + withGeminiImageGateAPIKey(c, tc.group) + + _, _ = svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), tc.model, tc.action, false, body) + require.Equal(t, 1, httpStub.calls) + require.NotEqual(t, http.StatusBadRequest, rec.Code) + }) + } +} + +func TestGeminiCompatForward_ImageGenerationDisabledGroupRejectsBeforeUpstream(t *testing.T) { + gin.SetMode(gin.TestMode) + svc, httpStub := newGeminiVertexAPIKeyTestService() + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + withGeminiImageGateAPIKey(c, imageDisabledGroup()) + body := []byte(`{"model":"gemini-3-pro-image","max_tokens":16,"messages":[{"role":"user","content":"draw a cat"}]}`) + + result, err := svc.Forward(context.Background(), c, geminiVertexAPIKeyTestAccount(""), body) + require.Error(t, err) + require.Nil(t, result) + require.Zero(t, httpStub.calls) + requireGeminiImagePermissionError(t, rec, "") +} + +func TestGeminiForwardAsChatCompletions_ImageGenerationDisabledGroupRejectsBeforeUpstream(t *testing.T) { + gin.SetMode(gin.TestMode) + svc, httpStub := newGeminiVertexAPIKeyTestService() + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := []byte(`{"model":"gemini-3.1-flash-image","messages":[{"role":"user","content":"draw a cat"}]}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + withGeminiImageGateAPIKey(c, imageDisabledGroup()) + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), body) + require.Error(t, err) + require.Nil(t, result) + require.Zero(t, httpStub.calls) + requireGeminiImagePermissionError(t, rec, "") +} + +func newAntigravityImageGateAccount() *Account { + account := newAntigravityCompatAccount(AccountTypeOAuth) + account.Credentials["model_mapping"] = map[string]any{ + "gemini-3-pro-image": "gemini-3-pro-image", + "my-draw": "gemini-3.1-flash-image", + } + return account +} + +func TestAntigravityForwards_ImageGenerationGate(t *testing.T) { + gin.SetMode(gin.TestMode) + geminiBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"draw a cat"}]}]}`) + claudeBody := []byte(`{"model":"my-draw","max_tokens":16,"stream":true,"messages":[{"role":"user","content":"draw a cat"}]}`) + chatBody := []byte(`{"model":"gemini-3-pro-image","stream":true,"messages":[{"role":"user","content":"draw a cat"}]}`) + + forwards := []struct { + name string + path string + body []byte + wantGoogleStatus string + forward func(*AntigravityGatewayService, *gin.Context, *Account) (*ForwardResult, error) + }{ + { + name: "ForwardGemini", + path: "/v1beta/models/gemini-3-pro-image:streamGenerateContent", + body: geminiBody, + wantGoogleStatus: "INVALID_ARGUMENT", + forward: func(svc *AntigravityGatewayService, c *gin.Context, account *Account) (*ForwardResult, error) { + return svc.ForwardGemini(context.Background(), c, account, "gemini-3-pro-image", "streamGenerateContent", true, geminiBody, false) + }, + }, + { + name: "Forward", + path: "/v1/messages", + body: claudeBody, + forward: func(svc *AntigravityGatewayService, c *gin.Context, account *Account) (*ForwardResult, error) { + return svc.Forward(context.Background(), c, account, claudeBody, false) + }, + }, + { + name: "ForwardAsChatCompletions", + path: "/v1/chat/completions", + body: chatBody, + forward: func(svc *AntigravityGatewayService, c *gin.Context, account *Account) (*ForwardResult, error) { + return svc.ForwardAsChatCompletions(context.Background(), c, account, chatBody, nil) + }, + }, + } + for _, fw := range forwards { + t.Run(fw.name+"/disabled group rejects", func(t *testing.T) { + upstream := &queuedHTTPUpstreamStub{} + svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, upstream) + c, rec := newAntigravityCompatContext(http.MethodPost, fw.path, fw.body) + withGeminiImageGateAPIKey(c, &Group{ID: 20, Platform: PlatformAntigravity, AllowImageGeneration: false}) + + result, err := fw.forward(svc, c, newAntigravityImageGateAccount()) + require.Error(t, err) + require.Nil(t, result) + require.Zero(t, upstream.callCount) + requireGeminiImagePermissionError(t, rec, fw.wantGoogleStatus) + }) + t.Run(fw.name+"/enabled group reaches upstream", func(t *testing.T) { + upstream := &queuedHTTPUpstreamStub{responses: []*http.Response{antigravityCompatSuccessResponse()}} + svc := newAntigravityCompatService(config.GatewayConfig{MaxLineSize: defaultMaxLineSize}, upstream) + c, rec := newAntigravityCompatContext(http.MethodPost, fw.path, fw.body) + withGeminiImageGateAPIKey(c, &Group{ID: 21, Platform: PlatformAntigravity, AllowImageGeneration: true}) + + _, _ = fw.forward(svc, c, newAntigravityImageGateAccount()) + require.Equal(t, 1, upstream.callCount) + require.NotEqual(t, http.StatusBadRequest, rec.Code) + }) + } +} diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index bafa5194f9b3..9ef4f29d079c 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -556,6 +556,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex if err != nil { return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", err.Error()) } + if geminiImageGenerationDenied(c, geminiReq, originalModel, mappedModel) { + return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", ImageGenerationPermissionMessage()) + } geminiReq = ensureGeminiFunctionCallThoughtSignatures(geminiReq) includeThoughts := req.Stream && claudeRequestEnablesThinking(body) geminiReq = applyGeminiCompatThinkingConfig(geminiReq, body, mappedModel, includeThoughts) @@ -1018,6 +1021,8 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex var usage *ClaudeUsage var firstTokenMs *int + clientDisconnect := false + streamFinished := false if req.Stream { streamRes, err := s.handleStreamingResponse(c, resp, startTime, originalModel) if err != nil { @@ -1025,6 +1030,8 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex } usage = streamRes.usage firstTokenMs = streamRes.firstTokenMs + clientDisconnect = streamRes.clientDisconnect + streamFinished = streamRes.finished } else { if useUpstreamStream { collected, usageObj, err := collectGeminiSSE(resp.Body, true) @@ -1052,6 +1059,10 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex imageInputSize := s.extractImageInputSize(body) imageSize := normalizeOpenAIImageSizeTier(imageInputSize) imageCount := resolveGeminiImageCount(c, originalModel, mappedModel) + // 客户端在收到结束分片前断开时只计已送达的图片：Claude 兼容流不向客户端转发图片。 + if clientDisconnect && !streamFinished { + imageCount = 0 + } return &ForwardResult{ RequestID: requestID, @@ -1065,6 +1076,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex Duration: time.Since(startTime), FirstTokenMs: firstTokenMs, ReasoningEffort: extractGeminiReasoningEffortFromBody(geminiReq), + ClientDisconnect: clientDisconnect, ImageCount: imageCount, ImageSize: imageSize, ImageInputSize: imageInputSize, @@ -1117,6 +1129,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. if account.Type == AccountTypeAPIKey || account.Type == AccountTypeServiceAccount { mappedModel = account.GetMappedModel(originalModel) } + if action != "countTokens" && geminiImageGenerationDenied(c, body, originalModel, mappedModel) { + return nil, s.writeGoogleError(c, http.StatusBadRequest, ImageGenerationPermissionMessage()) + } proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { @@ -1536,6 +1551,8 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. var usage *ClaudeUsage var firstTokenMs *int + clientDisconnect := false + streamFinished := false if stream { streamRes, err := s.handleNativeStreamingResponse(c, resp, startTime, isOAuth, account, requestID) @@ -1544,6 +1561,8 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. } usage = streamRes.usage firstTokenMs = streamRes.firstTokenMs + clientDisconnect = streamRes.clientDisconnect + streamFinished = streamRes.finished } else { if useUpstreamStream { var best geminiResponseSignal @@ -1578,6 +1597,10 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. imageInputSize := s.extractImageInputSize(body) imageSize := normalizeOpenAIImageSizeTier(imageInputSize) imageCount := resolveGeminiImageCount(c, originalModel, mappedModel) + // 客户端在收到结束分片前断开时只计已送达的图片。 + if clientDisconnect && !streamFinished { + imageCount = observedGeminiImageOutputs(c) + } return &ForwardResult{ RequestID: requestID, @@ -1591,6 +1614,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. ReasoningEffort: extractGeminiReasoningEffortFromBody(body), Duration: time.Since(startTime), FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnect, ImageCount: imageCount, ImageSize: imageSize, ImageInputSize: imageInputSize, @@ -2074,8 +2098,10 @@ func mapGeminiStatusToClaudeErrorType(status string) string { } type geminiStreamResult struct { - usage *ClaudeUsage - firstTokenMs *int + usage *ClaudeUsage + firstTokenMs *int + clientDisconnect bool // 客户端在流式传输中途断开，结果为断开前已收集的部分 + finished bool // 断开前已收到上游 finishReason } func (s *GeminiMessagesCompatService) handleNonStreamingResponse(c *gin.Context, resp *http.Response, originalModel string) (*ClaudeUsage, error) { @@ -2157,6 +2183,15 @@ func (s *GeminiMessagesCompatService) handleStreamingResponse(c *gin.Context, re for { line, err := reader.ReadString('\n') if err != nil && !errors.Is(err, io.EOF) { + if geminiStreamAbortedWithBillableOutput(c, err, &usage) { + logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] Client disconnected during messages stream, billing collected usage: %v", err) + return &geminiStreamResult{ + usage: &usage, + firstTokenMs: firstTokenMs, + clientDisconnect: true, + finished: finishReason != "", + }, nil + } return nil, fmt.Errorf("stream read error: %w", err) } @@ -2666,8 +2701,10 @@ func mergeCollectedTextParts(response map[string]any, textParts []string) map[st } type geminiNativeStreamResult struct { - usage *ClaudeUsage - firstTokenMs *int + usage *ClaudeUsage + firstTokenMs *int + clientDisconnect bool // 客户端在流式传输中途断开，结果为断开前已收集的部分 + finished bool // 断开前已收到上游 finishReason } func isGeminiInsufficientScope(headers http.Header, body []byte) bool { @@ -2844,10 +2881,13 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte var best geminiResponseSignal sawDataEvent := false fallback := &geminiSSEFallbackBody{} + clientDisconnect := false + finished := false for { line, err := reader.ReadString('\n') - if len(line) > 0 { + // 读错误时 line 是未以换行结尾的残行，不透传也不参与用量统计；EOF 时的末行完整有效。 + if len(line) > 0 && (err == nil || errors.Is(err, io.EOF)) { trimmed := strings.TrimRight(line, "\r\n") if strings.HasPrefix(trimmed, "data:") { payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) @@ -2880,6 +2920,10 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte observer.ObserveGemini(rawBytes) observeGeminiImageOutputs(c, rawBytes) + if gjson.GetBytes(rawBytes, "candidates.0.finishReason").String() != "" { + finished = true + } + if firstTokenMs == nil { ms := int(time.Since(startTime).Milliseconds()) firstTokenMs = &ms @@ -2907,13 +2951,18 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte break } if err != nil { - return nil, err + if !geminiStreamAbortedWithBillableOutput(c, err, usage) { + return nil, err + } + clientDisconnect = true + logger.LegacyPrintf("service.gemini_messages_compat", "[GeminiAPI] Client disconnected during native stream (account=%d), billing collected usage: %v", account.ID, err) + break } } s.finalizeGeminiSSESignal(c, account, true, upstreamRequestID, best, sawDataEvent, fallback) - return &geminiNativeStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil + return &geminiNativeStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnect, finished: finished}, nil } // unwrapGeminiResponse 解包 Gemini OAuth 响应中的 response 字段 @@ -3042,9 +3091,10 @@ func extractGeminiUsage(data []byte) *ClaudeUsage { } // 注意：Gemini 的 promptTokenCount 包含 cachedContentTokenCount， - // 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去 + // 但 Claude 的 input_tokens 不包含 cache_read_input_tokens，需要减去。 + // prompt 被拦截时上游可能只报 cachedContentTokenCount，差值下限为 0。 return &ClaudeUsage{ - InputTokens: prompt - cached, + InputTokens: max(prompt-cached, 0), OutputTokens: cand + thoughts, CacheReadInputTokens: cached, ImageOutputTokens: imageTokens, @@ -3774,12 +3824,24 @@ func isClaudeWebSearchToolMap(tool map[string]any) bool { // cleanToolSchema 清理工具的 JSON Schema，移除 Gemini 不支持的字段 func cleanToolSchema(schema any) any { + return cleanToolSchemaNode(schema, true) +} + +func cleanToolSchemaNode(schema any, schemaNode bool) any { if schema == nil { return nil } switch v := schema.(type) { case map[string]any: + if !schemaNode { + cleaned := make(map[string]any, len(v)) + for key, value := range v { + cleaned[key] = cleanToolSchemaNode(value, false) + } + return cleaned + } + cleaned := make(map[string]any) for key, value := range v { // 跳过不支持的字段 @@ -3787,11 +3849,73 @@ func cleanToolSchema(schema any) any { key == "$defs" || key == "definitions" || key == "additionalProperties" || key == "patternProperties" || key == "minLength" || key == "maxLength" || key == "minItems" || key == "maxItems" || key == "exclusiveMinimum" || - key == "deprecated" { + key == "deprecated" || key == "examples" { continue } - // 递归清理嵌套对象 - cleaned[key] = cleanToolSchema(value) + + switch key { + case "properties", "dependentSchemas": + if propertySchemas, ok := value.(map[string]any); ok { + cleanedProperties := make(map[string]any, len(propertySchemas)) + for propertyName, propertySchema := range propertySchemas { + cleanedProperties[propertyName] = cleanToolSchemaNode(propertySchema, true) + } + cleaned[key] = cleanedProperties + } else { + cleaned[key] = cleanToolSchemaNode(value, false) + } + case "items", "contains", "propertyNames", "if", "then", "else", "not", "contentSchema", + "additionalItems", "unevaluatedItems", "unevaluatedProperties": + cleaned[key] = cleanToolSchemaNode(value, true) + case "dependencies": + if dependencies, ok := value.(map[string]any); ok { + cleanedDependencies := make(map[string]any, len(dependencies)) + for propertyName, dependency := range dependencies { + if dependencySchema, ok := dependency.(map[string]any); ok { + cleanedDependencies[propertyName] = cleanToolSchemaNode(dependencySchema, true) + } else { + cleanedDependencies[propertyName] = cleanToolSchemaNode(dependency, false) + } + } + cleaned[key] = cleanedDependencies + } else { + cleaned[key] = cleanToolSchemaNode(value, false) + } + case "anyOf", "allOf", "oneOf", "prefixItems": + if schemas, ok := value.([]any); ok { + cleanedSchemas := make([]any, len(schemas)) + for i, childSchema := range schemas { + cleanedSchemas[i] = cleanToolSchemaNode(childSchema, true) + } + cleaned[key] = cleanedSchemas + } else { + cleaned[key] = cleanToolSchemaNode(value, false) + } + default: + cleaned[key] = cleanToolSchemaNode(value, false) + } + } + if constValue, exists := v["const"]; exists { + if stringValue, ok := constValue.(string); ok { + if originalEnum, hasEnum := v["enum"].([]any); hasEnum { + intersection := make([]any, 0, 1) + for _, value := range originalEnum { + if enumString, ok := value.(string); ok && enumString == stringValue { + intersection = append(intersection, stringValue) + break + } + } + cleaned["enum"] = intersection + } else { + cleaned["enum"] = []any{stringValue} + } + if _, hasType := cleaned["type"]; !hasType { + cleaned["type"] = "STRING" + } + } + // Non-string constants are unsupported: omit the constant, retain the + // original type, and do not stringify the constant or infer a type. + delete(cleaned, "const") } if enum, ok := cleaned["enum"].([]any); ok { if normalized, ok := normalizeGeminiEnum(enum); ok { @@ -3826,7 +3950,7 @@ func cleanToolSchema(schema any) any { case []any: cleaned := make([]any, len(v)) for i, item := range v { - cleaned[i] = cleanToolSchema(item) + cleaned[i] = cleanToolSchemaNode(item, schemaNode) } return cleaned default: diff --git a/backend/internal/service/gemini_messages_compat_service_test.go b/backend/internal/service/gemini_messages_compat_service_test.go index d29e2c7d8b21..86f15770c137 100644 --- a/backend/internal/service/gemini_messages_compat_service_test.go +++ b/backend/internal/service/gemini_messages_compat_service_test.go @@ -509,6 +509,226 @@ func TestCleanToolSchema_DropsEnumWithNonScalarValue(t *testing.T) { require.NotContains(t, cleaned, "enum") } +func TestCleanToolSchema_PreservesSchemaKeywordNamedProperties(t *testing.T) { + schema := map[string]any{ + "type": "object", + "examples": []any{"root example"}, + "properties": map[string]any{ + "const": map[string]any{ + "type": "string", + "const": "property const", + }, + "enum": map[string]any{ + "type": "string", + "const": "property enum", + }, + "examples": map[string]any{ + "type": "string", + }, + "type": map[string]any{ + "type": "string", + }, + }, + } + + cleaned, ok := cleanToolSchema(schema).(map[string]any) + require.True(t, ok) + require.NotContains(t, cleaned, "examples") + + properties, ok := cleaned["properties"].(map[string]any) + require.True(t, ok) + require.Contains(t, properties, "const") + require.Contains(t, properties, "enum") + require.Contains(t, properties, "examples") + require.Contains(t, properties, "type") + + constSchema, ok := properties["const"].(map[string]any) + require.True(t, ok) + require.Equal(t, []any{"property const"}, constSchema["enum"]) + + enumPropertySchema, ok := properties["enum"].(map[string]any) + require.True(t, ok) + require.Equal(t, []any{"property enum"}, enumPropertySchema["enum"]) +} + +func TestCleanToolSchema_NormalizesConstAndExamplesAtNestedSchemaNodes(t *testing.T) { + schema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "values": map[string]any{ + "type": "array", + "items": map[string]any{ + "type": "string", + "examples": []any{"item example"}, + "anyOf": []any{ + map[string]any{"type": "string", "const": "a"}, + map[string]any{"type": "string", "examples": []any{"branch example"}}, + }, + }, + }, + }, + } + + cleaned, ok := cleanToolSchema(schema).(map[string]any) + require.True(t, ok) + properties, ok := cleaned["properties"].(map[string]any) + require.True(t, ok) + values, ok := properties["values"].(map[string]any) + require.True(t, ok) + items, ok := values["items"].(map[string]any) + require.True(t, ok) + require.NotContains(t, items, "examples") + anyOf, ok := items["anyOf"].([]any) + require.True(t, ok) + require.Len(t, anyOf, 2) + firstBranch, ok := anyOf[0].(map[string]any) + require.True(t, ok) + require.Equal(t, []any{"a"}, firstBranch["enum"]) + secondBranch, ok := anyOf[1].(map[string]any) + require.True(t, ok) + require.NotContains(t, secondBranch, "examples") +} + +func TestCleanToolSchema_TraversesRetainedSchemaValuedKeywords(t *testing.T) { + schema := map[string]any{ + "type": "object", + "additionalItems": map[string]any{ + "type": "string", + "const": "additional", + "examples": []any{"additional example"}, + }, + "unevaluatedItems": map[string]any{ + "type": "string", + "const": "unevaluated item", + "examples": []any{"item example"}, + }, + "unevaluatedProperties": map[string]any{ + "type": "string", + "const": "unevaluated property", + "examples": []any{"property example"}, + }, + "dependencies": map[string]any{ + "schema dependency": map[string]any{ + "type": "string", + "const": "dependency", + "examples": []any{"dependency example"}, + }, + "required dependency": []any{"first", "second"}, + }, + } + + cleaned, ok := cleanToolSchema(schema).(map[string]any) + require.True(t, ok) + expectedEnums := map[string][]any{ + "additionalItems": {"additional"}, + "unevaluatedItems": {"unevaluated item"}, + "unevaluatedProperties": {"unevaluated property"}, + } + for keyword, expectedEnum := range expectedEnums { + child, ok := cleaned[keyword].(map[string]any) + require.True(t, ok) + require.Equal(t, expectedEnum, child["enum"]) + require.NotContains(t, child, "const") + require.NotContains(t, child, "examples") + } + + dependencies, ok := cleaned["dependencies"].(map[string]any) + require.True(t, ok) + dependencySchema, ok := dependencies["schema dependency"].(map[string]any) + require.True(t, ok) + require.Equal(t, []any{"dependency"}, dependencySchema["enum"]) + require.NotContains(t, dependencySchema, "const") + require.NotContains(t, dependencySchema, "examples") + require.Equal(t, []any{"first", "second"}, dependencies["required dependency"]) +} + +func TestCleanToolSchema_IntersectsStringConstWithEnum(t *testing.T) { + tests := []struct { + name string + constValue string + enum []any + expected []any + }{ + {name: "matching", constValue: "browser", enum: []any{"browser", "shell"}, expected: []any{"browser"}}, + {name: "contradictory", constValue: "browser", enum: []any{"shell"}, expected: []any{}}, + {name: "numeric string collision", constValue: "7", enum: []any{float64(7)}, expected: []any{}}, + {name: "empty string", constValue: "", enum: []any{"", "shell"}, expected: []any{""}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + schema := map[string]any{ + "type": "string", + "const": test.constValue, + "enum": test.enum, + } + + cleaned, ok := cleanToolSchema(schema).(map[string]any) + require.True(t, ok) + require.NotContains(t, cleaned, "const") + require.Equal(t, test.expected, cleaned["enum"]) + }) + } + + cleaned, ok := cleanToolSchema(map[string]any{"const": "browser"}).(map[string]any) + require.True(t, ok) + require.Equal(t, "STRING", cleaned["type"]) + require.Equal(t, []any{"browser"}, cleaned["enum"]) +} + +func TestCleanToolSchema_DropsNonStringConstWithoutInferringType(t *testing.T) { + for _, constValue := range []any{true, float64(7), map[string]any{"value": "x"}} { + cleaned, ok := cleanToolSchema(map[string]any{"const": constValue}).(map[string]any) + require.True(t, ok) + require.NotContains(t, cleaned, "const") + require.NotContains(t, cleaned, "type") + require.NotContains(t, cleaned, "enum") + } +} + +func TestCleanToolSchema_DropsNonStringConstAndRetainsExistingTypeAndEnum(t *testing.T) { + schema := map[string]any{ + "type": "integer", + "const": float64(7), + "enum": []any{float64(8)}, + } + + cleaned, ok := cleanToolSchema(schema).(map[string]any) + require.True(t, ok) + require.NotContains(t, cleaned, "const") + require.Equal(t, "INTEGER", cleaned["type"]) + require.Equal(t, []any{"8"}, cleaned["enum"]) +} + +func TestCleanToolSchema_PreservesNestedOneOfCombinator(t *testing.T) { + schema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "value": map[string]any{ + "oneOf": []any{ + map[string]any{"const": "a"}, + map[string]any{"type": "string", "const": "b"}, + }, + }, + }, + } + + cleaned, ok := cleanToolSchema(schema).(map[string]any) + require.True(t, ok) + properties, ok := cleaned["properties"].(map[string]any) + require.True(t, ok) + valueSchema, ok := properties["value"].(map[string]any) + require.True(t, ok) + oneOf, ok := valueSchema["oneOf"].([]any) + require.True(t, ok) + require.Len(t, oneOf, 2) + firstBranch, ok := oneOf[0].(map[string]any) + require.True(t, ok) + secondBranch, ok := oneOf[1].(map[string]any) + require.True(t, ok) + require.Equal(t, []any{"a"}, firstBranch["enum"]) + require.Equal(t, []any{"b"}, secondBranch["enum"]) +} + func TestConvertClaudeToolsToGeminiTools_PreservesWebSearchAlongsideFunctions(t *testing.T) { tools := []any{ map[string]any{ @@ -947,6 +1167,16 @@ func TestExtractGeminiUsage(t *testing.T) { CacheReadInputTokens: 0, }, }, + { + name: "prompt 被拦截只报 cachedContentTokenCount — input 下限为 0", + input: `{"promptFeedback":{"blockReason":"PROHIBITED_CONTENT"},"usageMetadata":{"cachedContentTokenCount":12980}}`, + wantNil: false, + wantUsage: &ClaudeUsage{ + InputTokens: 0, + OutputTokens: 0, + CacheReadInputTokens: 12980, + }, + }, { name: "无 usageMetadata", input: `{"candidates":[]}`, diff --git a/backend/internal/service/gemini_native_response_signal_test.go b/backend/internal/service/gemini_native_response_signal_test.go index 78c061ff089a..1005eca1e7cb 100644 --- a/backend/internal/service/gemini_native_response_signal_test.go +++ b/backend/internal/service/gemini_native_response_signal_test.go @@ -180,6 +180,29 @@ func TestGeminiForwardNative_StreamPromptBlockedMarksContentPolicy(t *testing.T) require.Empty(t, upstreamErrorEventsFromContext(t, c)) } +// Vertex 拦截 prompt 时 usageMetadata 可能只带 cachedContentTokenCount、缺 promptTokenCount： +// input 下限为 0，缓存命中照上游所报计入 cache_read。 +func TestGeminiForwardNative_StreamPromptBlockedWithOnlyCachedTokensKeepsInputNonNegative(t *testing.T) { + gin.SetMode(gin.TestMode) + body := `data: {"promptFeedback":{"blockReason":"PROHIBITED_CONTENT"},"usageMetadata":{"cachedContentTokenCount":12980,"totalTokenCount":0,"trafficType":"ON_DEMAND"},"modelVersion":"gemini-3.5-flash-lite"} + +` + svc := newGeminiSignalService("text/event-stream", body) + c, rec := newGeminiNativeTestContext(t) + + result, err := svc.ForwardNative(context.Background(), c, geminiSignalTestAccount(), + "gemini-3.5-flash-lite", "streamGenerateContent", true, geminiSignalTestRequest()) + require.NoError(t, err) + require.Equal(t, body, rec.Body.String()) + require.Equal(t, 0, result.Usage.InputTokens) + require.Equal(t, 12980, result.Usage.CacheReadInputTokens) + require.Equal(t, 0, result.Usage.OutputTokens) + + streamErrs := GetOpsStreamErrors(c) + require.Len(t, streamErrs, 1) + require.Equal(t, "PROHIBITED_CONTENT", streamErrs[0].Code) +} + func TestGeminiForwardNative_EmptyStreamMarksUpstreamFailure(t *testing.T) { gin.SetMode(gin.TestMode) svc := newGeminiSignalService("text/event-stream", "") diff --git a/backend/internal/service/gemini_stream_client_abort.go b/backend/internal/service/gemini_stream_client_abort.go new file mode 100644 index 000000000000..d1dba1b8028e --- /dev/null +++ b/backend/internal/service/gemini_stream_client_abort.go @@ -0,0 +1,25 @@ +package service + +import ( + "context" + "errors" + "io" + + "github.com/gin-gonic/gin" +) + +// geminiStreamAbortedWithBillableOutput 判断流式读取中断是否由客户端断开引起，且断开前 +// 已观测到可计费的输出（token 用量或已回吐的图片）。满足时调用方以已收集的结果收尾并照常计费。 +// usageMetadata 为累计值，按断开前收到的最后一份计费：读到结束分片后断开即完整用量； +// 中途断开时中间分片可能只含输入计数或不含用量，此时按已有计数计费或不计费。 +func geminiStreamAbortedWithBillableOutput(c *gin.Context, err error, usage *ClaudeUsage) bool { + if err == nil || errors.Is(err, io.EOF) { + return false + } + clientGone := errors.Is(err, context.Canceled) || + (c != nil && c.Request != nil && c.Request.Context().Err() != nil) + if !clientGone { + return false + } + return usage.hasObservedTokens() || observedGeminiImageOutputs(c) > 0 +} diff --git a/backend/internal/service/gemini_stream_client_abort_test.go b/backend/internal/service/gemini_stream_client_abort_test.go new file mode 100644 index 000000000000..0825d5152aa5 --- /dev/null +++ b/backend/internal/service/gemini_stream_client_abort_test.go @@ -0,0 +1,244 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// geminiAbortingStreamBody 先吐出 data，读完后返回 abortErr；onAbort 用于模拟客户端断开时请求 context 被取消。 +type geminiAbortingStreamBody struct { + data *bytes.Reader + abortErr error + onAbort func() +} + +func (b *geminiAbortingStreamBody) Read(p []byte) (int, error) { + n, err := b.data.Read(p) + if err == io.EOF { + if b.onAbort != nil { + b.onAbort() + } + return n, b.abortErr + } + return n, err +} + +func (b *geminiAbortingStreamBody) Close() error { return nil } + +const ( + geminiAbortTextChunk = `data: {"candidates":[{"content":{"parts":[{"text":"hello"}],"role":"model"}}]}` + "\n\n" + geminiAbortFinalChunk = `data: {"candidates":[{"content":{"parts":[{"text":" world"}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}` + "\n\n" + geminiAbortImageChunk = `data: {"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"aGVsbG8="}}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1290,"totalTokenCount":1300}}` + "\n\n" + geminiAbortUsageChunk = `data: {"candidates":[{"content":{"parts":[{"text":"drawing"}],"role":"model"}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"totalTokenCount":12}}` + "\n\n" +) + +// newGeminiAbortingStreamService 返回的上游流在吐完 sse 后中断；clientGone=true 时中断伴随请求 context 取消。 +func newGeminiAbortingStreamService(t *testing.T, path string, reqBody []byte, sse string, clientGone bool) (*GeminiMessagesCompatService, *gin.Context, *httptest.ResponseRecorder) { + t.Helper() + gin.SetMode(gin.TestMode) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + reqCtx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + c.Request = httptest.NewRequest(http.MethodPost, path, bytes.NewReader(reqBody)).WithContext(reqCtx) + + body := &geminiAbortingStreamBody{data: bytes.NewReader([]byte(sse)), abortErr: io.ErrUnexpectedEOF} + if clientGone { + body.abortErr = context.Canceled + body.onAbort = cancel + } + svc, httpStub := newGeminiVertexAPIKeyTestService() + httpStub.response = &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: body, + } + return svc, c, rec +} + +func TestGeminiForwardNative_ClientDisconnectAfterFinalChunkBillsCollectedUsage(t *testing.T) { + reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`) + svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-flash-preview:streamGenerateContent", reqBody, geminiAbortTextChunk+geminiAbortFinalChunk, true) + + result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-flash-preview", "streamGenerateContent", true, reqBody) + require.NoError(t, err) + require.NotNil(t, result) + require.True(t, result.ClientDisconnect) + require.Equal(t, 10, result.Usage.InputTokens) + require.Equal(t, 5, result.Usage.OutputTokens) + require.Zero(t, result.ImageCount) +} + +func TestGeminiForwardNative_ClientDisconnectWithoutUsageStaysUnbilled(t *testing.T) { + reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`) + svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-flash-preview:streamGenerateContent", reqBody, geminiAbortTextChunk, true) + + result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-flash-preview", "streamGenerateContent", true, reqBody) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, result) +} + +func TestGeminiForwardNative_UpstreamReadErrorWithLiveClientStaysError(t *testing.T) { + reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`) + svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-flash-preview:streamGenerateContent", reqBody, geminiAbortTextChunk+geminiAbortFinalChunk, false) + + result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-flash-preview", "streamGenerateContent", true, reqBody) + require.ErrorIs(t, err, io.ErrUnexpectedEOF) + require.Nil(t, result) +} + +func TestGeminiForwardNative_ClientDisconnectBillsOnlyDeliveredImages(t *testing.T) { + reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"draw a cat"}]}]}`) + + t.Run("image delivered before disconnect", func(t *testing.T) { + svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-pro-image:streamGenerateContent", reqBody, geminiAbortImageChunk, true) + result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-pro-image", "streamGenerateContent", true, reqBody) + require.NoError(t, err) + require.True(t, result.ClientDisconnect) + require.Equal(t, 1, result.ImageCount) + }) + + t.Run("disconnect before the image arrives", func(t *testing.T) { + svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-pro-image:streamGenerateContent", reqBody, geminiAbortUsageChunk, true) + result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-pro-image", "streamGenerateContent", true, reqBody) + require.NoError(t, err) + require.True(t, result.ClientDisconnect) + require.Zero(t, result.ImageCount) + require.Equal(t, 2, result.Usage.OutputTokens) + }) +} + +func TestGeminiCompatForward_ClientDisconnectAfterFinalChunkBillsCollectedUsage(t *testing.T) { + reqBody := []byte(`{"model":"gemini-3-flash-preview","max_tokens":16,"stream":true,"messages":[{"role":"user","content":"hi"}]}`) + + t.Run("client gone", func(t *testing.T) { + svc, c, _ := newGeminiAbortingStreamService(t, "/v1/messages", reqBody, geminiAbortTextChunk+geminiAbortFinalChunk, true) + result, err := svc.Forward(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody) + require.NoError(t, err) + require.True(t, result.ClientDisconnect) + require.Equal(t, 10, result.Usage.InputTokens) + require.Equal(t, 5, result.Usage.OutputTokens) + }) + + t.Run("upstream read error with live client", func(t *testing.T) { + svc, c, _ := newGeminiAbortingStreamService(t, "/v1/messages", reqBody, geminiAbortTextChunk+geminiAbortFinalChunk, false) + result, err := svc.Forward(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody) + require.ErrorIs(t, err, io.ErrUnexpectedEOF) + require.Nil(t, result) + }) +} + +func TestGeminiForwardAsChatCompletions_ClientDisconnectAfterFinalChunkBillsCollectedUsage(t *testing.T) { + reqBody := []byte(`{"model":"gemini-3-flash-preview","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + + t.Run("client gone", func(t *testing.T) { + svc, c, _ := newGeminiAbortingStreamService(t, "/v1/chat/completions", reqBody, geminiAbortTextChunk+geminiAbortFinalChunk, true) + result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody) + require.NoError(t, err) + require.True(t, result.ClientDisconnect) + require.Equal(t, 10, result.Usage.InputTokens) + require.Equal(t, 5, result.Usage.OutputTokens) + }) + + t.Run("client gone before usage", func(t *testing.T) { + svc, c, _ := newGeminiAbortingStreamService(t, "/v1/chat/completions", reqBody, geminiAbortTextChunk, true) + result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, result) + }) +} + +func TestGeminiForwardNative_ClientDisconnectAfterFinishKeepsImageFallback(t *testing.T) { + // 结束分片已到达但图片不以 inlineData 回吐时，断开与正常完成同口径按模型名兜底。 + finishWithoutInline := `data: {"candidates":[{"content":{"parts":[{"fileData":{"mimeType":"image/png","fileUri":"gs://bucket/cat.png"}}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1290,"totalTokenCount":1300}}` + "\n\n" + reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"draw a cat"}]}]}`) + svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-pro-image:streamGenerateContent", reqBody, finishWithoutInline, true) + + result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-pro-image", "streamGenerateContent", true, reqBody) + require.NoError(t, err) + require.True(t, result.ClientDisconnect) + require.Equal(t, 1, result.ImageCount) +} + +func TestGeminiForwardNative_ClientDisconnectIgnoresTruncatedTrailingLine(t *testing.T) { + complete := `data: {"candidates":[{"content":{"parts":[{"text":"hello"}],"role":"model"}}],"usageMetadata":{"promptTokenCount":1000,"candidatesTokenCount":300,"totalTokenCount":1300}}` + "\n\n" + truncated := `data: {"candidates":[{"content":{"parts":[{"text":" more"}],"role":"model"}}],"usageMetadata":{` + reqBody := []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`) + svc, c, _ := newGeminiAbortingStreamService(t, "/v1beta/models/gemini-3-flash-preview:streamGenerateContent", reqBody, complete+truncated, true) + + result, err := svc.ForwardNative(context.Background(), c, geminiVertexAPIKeyTestAccount(""), "gemini-3-flash-preview", "streamGenerateContent", true, reqBody) + require.NoError(t, err) + require.True(t, result.ClientDisconnect) + require.Equal(t, 1000, result.Usage.InputTokens) + require.Equal(t, 300, result.Usage.OutputTokens) +} + +func TestGeminiForwardAsChatCompletions_ClientDisconnectImageCount(t *testing.T) { + t.Run("disconnect before finish bills no image", func(t *testing.T) { + reqBody := []byte(`{"model":"gemini-3-pro-image","stream":true,"messages":[{"role":"user","content":"draw a cat"}]}`) + svc, c, _ := newGeminiAbortingStreamService(t, "/v1/chat/completions", reqBody, geminiAbortUsageChunk, true) + result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody) + require.NoError(t, err) + require.True(t, result.ClientDisconnect) + require.Zero(t, result.ImageCount) + require.Equal(t, 2, result.Usage.OutputTokens) + }) + + t.Run("disconnect after finish keeps completed-request image count", func(t *testing.T) { + reqBody := []byte(`{"model":"gemini-3-pro-image","stream":true,"messages":[{"role":"user","content":"draw a cat"}]}`) + svc, c, _ := newGeminiAbortingStreamService(t, "/v1/chat/completions", reqBody, geminiAbortImageChunk, true) + result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody) + require.NoError(t, err) + require.True(t, result.ClientDisconnect) + require.Equal(t, 1, result.ImageCount) + }) +} + +// geminiFailingResponseWriter 模拟客户端已断开：所有写入失败。 +type geminiFailingResponseWriter struct { + header http.Header +} + +func (w *geminiFailingResponseWriter) Header() http.Header { return w.header } +func (w *geminiFailingResponseWriter) Write([]byte) (int, error) { return 0, io.ErrClosedPipe } +func (w *geminiFailingResponseWriter) WriteHeader(int) {} +func (w *geminiFailingResponseWriter) Flush() {} + +func TestGeminiForwardAsChatCompletions_WriteFailureOnImageModelBillsNoImage(t *testing.T) { + gin.SetMode(gin.TestMode) + reqBody := []byte(`{"model":"gemini-3-pro-image","stream":true,"messages":[{"role":"user","content":"draw a cat"}]}`) + c, _ := gin.CreateTestContext(&geminiFailingResponseWriter{header: http.Header{}}) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(reqBody)) + svc, httpStub := newGeminiVertexAPIKeyTestService() + httpStub.response = &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(bytes.NewReader([]byte(geminiAbortUsageChunk))), + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody) + require.NoError(t, err) + require.True(t, result.ClientDisconnect) + require.Zero(t, result.ImageCount) +} + +func TestGeminiCompatForward_ClientDisconnectBeforeFinishBillsNoImage(t *testing.T) { + inlineWithoutFinish := `data: {"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"aGVsbG8="}}],"role":"model"}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":1290,"totalTokenCount":1300}}` + "\n\n" + reqBody := []byte(`{"model":"gemini-3-pro-image","max_tokens":16,"stream":true,"messages":[{"role":"user","content":"draw a cat"}]}`) + svc, c, _ := newGeminiAbortingStreamService(t, "/v1/messages", reqBody, inlineWithoutFinish, true) + + result, err := svc.Forward(context.Background(), c, geminiVertexAPIKeyTestAccount(""), reqBody) + require.NoError(t, err) + require.True(t, result.ClientDisconnect) + require.Zero(t, result.ImageCount) + require.Equal(t, 1290, result.Usage.OutputTokens) +} diff --git a/backend/internal/service/image_storage.go b/backend/internal/service/image_storage.go index b160846b4716..f0e0d110d948 100644 --- a/backend/internal/service/image_storage.go +++ b/backend/internal/service/image_storage.go @@ -8,10 +8,16 @@ import ( "fmt" "io" "mime" + "net" "net/http" + "net/url" "strconv" "strings" + "sync" + "syscall" "time" + + "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" ) const defaultImageMaxDownloadBytes int64 = 32 << 20 // 32 MiB @@ -53,8 +59,123 @@ func NewImageResultUploader(storage ImageStorage, prefix string, maxDownloadByte } } -func defaultImageDownloadHTTPClient() *http.Client { - return &http.Client{Timeout: 60 * time.Second} +// defaultImageDownloadHTTPClient 是进程内共享的图片下载客户端，复用连接池。 +var defaultImageDownloadHTTPClient = sync.OnceValue(func() *http.Client { + return newImageDownloadHTTPClient(http.ProxyFromEnvironment) +}) + +// newImageDownloadHTTPClient 返回只允许下载公网目的地的客户端。 +// 图片 url 来自上游响应，是不可信输入：目的地为 localhost、回环、私网、链路本地、 +// 云元数据等地址时一律拒绝，重定向的每一跳都经同一 RoundTrip 校验。 +// 不经代理时在 socket 层校验真实连接 IP（防 DNS rebinding）；经环境代理时由代理解析目标， +// 只能在发出前校验本地解析结果。 +func newImageDownloadHTTPClient(proxy func(*http.Request) (*url.URL, error)) *http.Client { + direct := newImageDownloadBaseTransport(imageDownloadDialControl) + proxied := newImageDownloadBaseTransport(nil) + proxied.Proxy = proxy + return &http.Client{ + Timeout: 60 * time.Second, + Transport: &imageDownloadTransport{direct: direct, proxied: proxied, proxy: proxy, resolvesPrivate: isPrivateOrLoopbackHost}, + } +} + +// newImageDownloadBaseTransport 返回与 http.DefaultTransport 参数一致、但不带代理的 Transport。 +// control 非空时在每个 socket 连接前校验目标地址。 +func newImageDownloadBaseTransport(control func(network, address string, c syscall.RawConn) error) *http.Transport { + dialer := &net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second, Control: control} + return &http.Transport{ + DialContext: dialer.DialContext, + ForceAttemptHTTP2: true, + MaxIdleConns: 100, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 10 * time.Second, + ExpectContinueTimeout: 1 * time.Second, + } +} + +// imageDownloadDialControl 拒绝连接到非公网 IP。address 是 DNS 解析后的真实连接地址， +// 多个解析结果时 Dialer 会跳过被拒地址继续尝试下一个。 +func imageDownloadDialControl(_, address string, _ syscall.RawConn) error { + host, _, err := net.SplitHostPort(address) + if err != nil { + return err + } + if ip := net.ParseIP(stripIPv6Zone(host)); ip == nil || isPrivateIP(ip) { + return &net.AddrError{Err: "address not allowed for image download", Addr: address} + } + return nil +} + +// imageDownloadTransport 按请求是否走环境代理分派到对应 Transport，并在发出前拒绝非公网目的地。 +type imageDownloadTransport struct { + direct http.RoundTripper + proxied http.RoundTripper + proxy func(*http.Request) (*url.URL, error) + // resolvesPrivate 报告 host 的本地解析结果是否含非公网地址，仅用于经代理的请求。 + resolvesPrivate func(ctx context.Context, host string) (bool, error) +} + +func (t *imageDownloadTransport) RoundTrip(req *http.Request) (*http.Response, error) { + next, err := t.route(req) + if err != nil { + if req.Body != nil { + _ = req.Body.Close() + } + return nil, err + } + return next.RoundTrip(req) +} + +func (t *imageDownloadTransport) route(req *http.Request) (http.RoundTripper, error) { + if req.URL == nil { + return nil, errors.New("image url is empty") + } + host := req.URL.Hostname() + if isBlockedImageDownloadHost(host) { + return nil, fmt.Errorf("image url host is not allowed: %s", host) + } + var proxyURL *url.URL + if t.proxy != nil { + var err error + if proxyURL, err = t.proxy(req); err != nil { + return nil, err + } + } + if proxyURL == nil { + return t.direct, nil + } + if t.resolvesPrivate == nil { + return nil, errors.New("image download proxy host check is not configured") + } + blocked, err := t.resolvesPrivate(req.Context(), host) + if err != nil { + return nil, fmt.Errorf("resolve image url host: %w", err) + } + if blocked { + return nil, fmt.Errorf("image url host is not allowed: %s", host) + } + return t.proxied, nil +} + +// isBlockedImageDownloadHost 报告 host 是否为空、被拒绝的主机名，或非公网 IP 字面量。 +// 域名的解析结果由 socket 层或代理分支再校验。 +func isBlockedImageDownloadHost(host string) bool { + host = strings.TrimSuffix(strings.ToLower(strings.TrimSpace(host)), ".") + if urlvalidator.IsBlockedHost(host) || isBlockedHostname(host) { + return true + } + if ip := net.ParseIP(stripIPv6Zone(host)); ip != nil { + return isPrivateIP(ip) + } + return false +} + +// stripIPv6Zone 去掉 IPv6 字面量的 zone（如 fe80::1%eth0），使其可被 net.ParseIP 解析。 +func stripIPv6Zone(host string) string { + if i := strings.LastIndexByte(host, '%'); i >= 0 { + return host[:i] + } + return host } // Rewrite 将 result（上游生图响应 JSON）里的每张图片转存到对象存储， @@ -215,9 +336,12 @@ func (u *ImageResultUploader) download(ctx context.Context, rawURL string) ([]by if int64(len(data)) > limit { return nil, "", fmt.Errorf("downloaded image exceeds %d bytes", limit) } - contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0]) + contentType := strings.ToLower(strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])) if !strings.HasPrefix(contentType, "image/") { - contentType = detectImageContentType(data) + // 声明与嗅探都不是图片时拒绝，避免把任意响应体（内部接口的 JSON/HTML 等）转存给租户。 + if contentType = detectedImageContentType(data); contentType == "" { + return nil, "", errors.New("download image: content is not an image") + } } return data, contentType, nil } diff --git a/backend/internal/service/image_storage_test.go b/backend/internal/service/image_storage_test.go index 3a320671ee12..960b95857969 100644 --- a/backend/internal/service/image_storage_test.go +++ b/backend/internal/service/image_storage_test.go @@ -1,12 +1,17 @@ package service import ( + "bytes" "context" "encoding/base64" "encoding/json" "errors" + "io" + "net" "net/http" "net/http/httptest" + "net/url" + "sync/atomic" "testing" "time" @@ -73,18 +78,24 @@ func TestImageResultUploaderRewritesB64JSON(t *testing.T) { } func TestImageResultUploaderRewritesURL(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", "image/png") - _, _ = w.Write(pngBytes) - })) - defer upstream.Close() + var requested string + client := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + requested = req.URL.String() + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"image/png"}}, + Body: io.NopCloser(bytes.NewReader(pngBytes)), + Request: req, + }, nil + })} storage := &fakeImageStorage{} - uploader := NewImageResultUploader(storage, "images/", 0, nil) + uploader := NewImageResultUploader(storage, "images/", 0, client) - result := json.RawMessage(`{"created":1,"data":[{"url":"` + upstream.URL + `/pic.png"}]}`) + result := json.RawMessage(`{"created":1,"data":[{"url":"https://images.example.test/pic.png"}]}`) out, err := uploader.Rewrite(context.Background(), "imgtask_xyz", result) require.NoError(t, err) + require.Equal(t, "https://images.example.test/pic.png", requested) require.Len(t, storage.saved, 1) require.Equal(t, pngBytes, storage.saved[0].data) @@ -244,3 +255,243 @@ func TestImageTaskServiceCompleteOffloadFailureMarksFailed(t *testing.T) { require.Contains(t, string(got.Error), "object storage") require.NotContains(t, string(got.Result), "b64_json", "failed offload must not persist base64 to Redis") } + +func TestImageResultUploaderDefaultClientRejectsPrivateURL(t *testing.T) { + var hits atomic.Int32 + internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits.Add(1) + _, _ = w.Write([]byte("GIF89a-payload")) + })) + defer internal.Close() + + storage := &fakeImageStorage{} + uploader := NewImageResultUploader(storage, "images/", 0, nil) + result := json.RawMessage(`{"data":[{"url":"` + internal.URL + `/a.gif"}]}`) + + _, err := uploader.Rewrite(context.Background(), "imgtask_private", result) + require.ErrorContains(t, err, "image url host is not allowed") + require.Zero(t, hits.Load(), "internal service must not be contacted") + require.Empty(t, storage.saved) +} + +func TestImageTaskServiceCompletePrivateURLMarksFailed(t *testing.T) { + var hits atomic.Int32 + internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits.Add(1) + _, _ = w.Write(pngBytes) + })) + defer internal.Close() + + store := &imageTaskMemoryStore{} + storage := &fakeImageStorage{} + svc := NewImageTaskServiceWithUploader(store, NewImageResultUploader(storage, "images/", 0, nil), time.Hour, time.Minute) + owner := ImageTaskOwner{UserID: 1, APIKeyID: 2} + created, err := svc.Create(context.Background(), owner) + require.NoError(t, err) + + result := json.RawMessage(`{"data":[{"url":"` + internal.URL + `/a.gif"}]}`) + require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, result)) + + got, err := svc.Get(context.Background(), owner, created.ID) + require.NoError(t, err) + require.Equal(t, ImageTaskStatusFailed, got.Status) + require.Empty(t, got.ImageURL) + require.NotContains(t, string(got.Result), internal.URL) + require.Zero(t, hits.Load()) + require.Empty(t, storage.saved) +} + +func newCountingImageDownloadTransport(proxyURL *url.URL) (*imageDownloadTransport, *atomic.Int32, *atomic.Int32) { + var directCalls, proxiedCalls atomic.Int32 + respond := func(counter *atomic.Int32) http.RoundTripper { + return roundTripFunc(func(req *http.Request) (*http.Response, error) { + counter.Add(1) + return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewReader(pngBytes)), Request: req}, nil + }) + } + return &imageDownloadTransport{ + direct: respond(&directCalls), + proxied: respond(&proxiedCalls), + proxy: func(*http.Request) (*url.URL, error) { return proxyURL, nil }, + resolvesPrivate: func(ctx context.Context, host string) (bool, error) { + if host == "intranet.example.test" { + return true, nil + } + if host == "unresolvable.example.test" { + return false, errors.New("no such host") + } + return isPrivateOrLoopbackHost(ctx, host) + }, + }, &directCalls, &proxiedCalls +} + +func TestImageDownloadTransportRejectsNonPublicHosts(t *testing.T) { + hosts := []string{ + "http://127.0.0.1:18082/a.gif", + "http://localhost/a.png", + "http://LOCALHOST./a.png", + "http://foo.localhost/a.png", + "http://10.0.0.8/a.png", + "http://172.16.3.4/a.png", + "http://192.168.1.1/a.png", + "http://100.64.0.1/a.png", + "http://169.254.169.254/latest/meta-data/", + "http://metadata.google.internal/computeMetadata/v1/", + "http://0.0.0.0/a.png", + "http://[::1]/a.png", + "http://[::ffff:127.0.0.1]/a.png", + "http://[fd00::1]/a.png", + "http://[fe80::1%25eth0]/a.png", + "http:///a.png", + } + for _, proxied := range []bool{false, true} { + for _, rawURL := range hosts { + var proxyURL *url.URL + if proxied { + proxyURL = &url.URL{Scheme: "http", Host: "127.0.0.1:7890"} + } + transport, directCalls, proxiedCalls := newCountingImageDownloadTransport(proxyURL) + req, err := http.NewRequest(http.MethodGet, rawURL, nil) + require.NoError(t, err) + + _, err = transport.RoundTrip(req) + require.Error(t, err, "proxied=%v url=%s", proxied, rawURL) + require.Zero(t, directCalls.Load(), rawURL) + require.Zero(t, proxiedCalls.Load(), rawURL) + } + } +} + +func TestImageDownloadTransportRoutesPublicHosts(t *testing.T) { + transport, directCalls, proxiedCalls := newCountingImageDownloadTransport(nil) + req, err := http.NewRequest(http.MethodGet, "https://203.0.113.10/a.png", nil) + require.NoError(t, err) + resp, err := transport.RoundTrip(req) + require.NoError(t, err) + _ = resp.Body.Close() + require.EqualValues(t, 1, directCalls.Load()) + require.Zero(t, proxiedCalls.Load()) + + transport, directCalls, proxiedCalls = newCountingImageDownloadTransport(&url.URL{Scheme: "http", Host: "127.0.0.1:7890"}) + req, err = http.NewRequest(http.MethodGet, "https://203.0.113.10/a.png", nil) + require.NoError(t, err) + resp, err = transport.RoundTrip(req) + require.NoError(t, err) + _ = resp.Body.Close() + require.Zero(t, directCalls.Load()) + require.EqualValues(t, 1, proxiedCalls.Load(), "a loopback env proxy must still be usable for public targets") +} + +func TestImageDownloadClientRejectsRedirectToPrivateHost(t *testing.T) { + var calls atomic.Int32 + transport := &imageDownloadTransport{ + direct: roundTripFunc(func(req *http.Request) (*http.Response, error) { + calls.Add(1) + return &http.Response{ + StatusCode: http.StatusFound, + Header: http.Header{"Location": []string{"http://169.254.169.254/latest/meta-data/"}}, + Body: io.NopCloser(bytes.NewReader(nil)), + Request: req, + }, nil + }), + } + uploader := NewImageResultUploader(&fakeImageStorage{}, "images/", 0, &http.Client{Transport: transport}) + + _, err := uploader.Rewrite(context.Background(), "imgtask_redirect", json.RawMessage(`{"data":[{"url":"https://images.example.test/a.png"}]}`)) + require.ErrorContains(t, err, "image url host is not allowed: 169.254.169.254") + require.EqualValues(t, 1, calls.Load(), "the private redirect hop must not be sent") +} + +func TestImageDownloadClientDialsOnlyPublicAddresses(t *testing.T) { + var hits atomic.Int32 + internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits.Add(1) + _, _ = w.Write(pngBytes) + })) + defer internal.Close() + + client := newImageDownloadHTTPClient(nil) + direct, ok := client.Transport.(*imageDownloadTransport).direct.(*http.Transport) + require.True(t, ok) + require.Nil(t, direct.Proxy, "the direct transport must not pick up an environment proxy") + + // 绕过 RoundTrip 预检直接 dial：IP 字面量与经 DNS 解析到回环的域名都必须在 socket 层被拒。 + conn, err := direct.DialContext(context.Background(), "tcp", internal.Listener.Addr().String()) + if conn != nil { + _ = conn.Close() + } + require.ErrorContains(t, err, "address not allowed for image download") + _, port, err := net.SplitHostPort(internal.Listener.Addr().String()) + require.NoError(t, err) + conn, err = direct.DialContext(context.Background(), "tcp", net.JoinHostPort("localhost", port)) + if conn != nil { + _ = conn.Close() + } + require.ErrorContains(t, err, "address not allowed for image download") + require.Zero(t, hits.Load()) +} + +func newStaticImageDownloadClient(contentType string, body []byte) *http.Client { + return &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + header := http.Header{} + if contentType != "" { + header.Set("Content-Type", contentType) + } + return &http.Response{StatusCode: http.StatusOK, Header: header, Body: io.NopCloser(bytes.NewReader(body)), Request: req}, nil + })} +} + +func TestImageResultUploaderDownloadContentTypes(t *testing.T) { + avif := []byte("\x00\x00\x00\x1cftypavif-not-sniffable") + tests := []struct { + name string + contentType string + body []byte + wantType string + wantErr string + }{ + {name: "declared png", contentType: "image/png", body: pngBytes, wantType: "image/png"}, + {name: "octet-stream sniffed png", contentType: "application/octet-stream", body: pngBytes, wantType: "image/png"}, + {name: "missing header sniffed png", body: pngBytes, wantType: "image/png"}, + {name: "declared avif not sniffable", contentType: "Image/AVIF; charset=binary", body: avif, wantType: "image/avif"}, + {name: "json body", contentType: "application/json", body: []byte(`{"ok":true}`), wantErr: "content is not an image"}, + {name: "html body", contentType: "text/html", body: []byte("admin"), wantErr: "content is not an image"}, + {name: "empty body", body: nil, wantErr: "content is not an image"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + storage := &fakeImageStorage{} + uploader := NewImageResultUploader(storage, "images/", 0, newStaticImageDownloadClient(tt.contentType, tt.body)) + _, err := uploader.Rewrite(context.Background(), "imgtask_ct", json.RawMessage(`{"data":[{"url":"https://images.example.test/a"}]}`)) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + require.Empty(t, storage.saved) + return + } + require.NoError(t, err) + require.Len(t, storage.saved, 1) + require.Equal(t, tt.wantType, storage.saved[0].contentType) + }) + } +} + +func TestImageDownloadDialControl(t *testing.T) { + for _, address := range []string{"127.0.0.1:443", "10.1.2.3:80", "100.100.100.200:80", "169.254.169.254:80", "[::1]:443", "[fe80::1%eth0]:443", "[::ffff:192.168.0.1]:443", "bad"} { + require.Error(t, imageDownloadDialControl("tcp", address, nil), address) + } + for _, address := range []string{"203.0.113.10:443", "[2001:db8::1]:443"} { + require.NoError(t, imageDownloadDialControl("tcp", address, nil), address) + } +} + +func TestImageDownloadTransportProxiedChecksResolvedHost(t *testing.T) { + for _, host := range []string{"intranet.example.test", "unresolvable.example.test"} { + transport, directCalls, proxiedCalls := newCountingImageDownloadTransport(&url.URL{Scheme: "http", Host: "127.0.0.1:7890"}) + req, err := http.NewRequest(http.MethodGet, "https://"+host+"/a.png", nil) + require.NoError(t, err) + _, err = transport.RoundTrip(req) + require.Error(t, err, host) + require.Zero(t, directCalls.Load(), host) + require.Zero(t, proxiedCalls.Load(), "a proxied request must fail closed when the target is private or unresolvable: %s", host) + } +} diff --git a/backend/internal/service/model_rate_limit.go b/backend/internal/service/model_rate_limit.go index 8245362255be..f0c9d2967a3c 100644 --- a/backend/internal/service/model_rate_limit.go +++ b/backend/internal/service/model_rate_limit.go @@ -2,6 +2,7 @@ package service import ( "context" + "maps" "strings" "time" @@ -246,18 +247,13 @@ func (a *Account) modelRateLimitReason(scope string) string { return reason } +// setAccountModelRateLimitSnapshot 把模型级限流写进本账号对象的 Extra 快照。 +// 写时复制：Extra 与其中的 model_rate_limits 可能与调度快照母本及其它副本共享， +// 这里只读取旧 map，写入新复制的 map 后替换 account.Extra，不修改任何已有 map。 func setAccountModelRateLimitSnapshot(account *Account, scope string, resetAt time.Time, reason string, now time.Time) { if account == nil || strings.TrimSpace(scope) == "" { return } - if account.Extra == nil { - account.Extra = make(map[string]any) - } - limits, ok := account.Extra[modelRateLimitsKey].(map[string]any) - if !ok { - limits = make(map[string]any) - account.Extra[modelRateLimitsKey] = limits - } payload := map[string]any{ "rate_limited_at": now.UTC().Format(time.RFC3339), "rate_limit_reset_at": resetAt.UTC().Format(time.RFC3339), @@ -265,5 +261,14 @@ func setAccountModelRateLimitSnapshot(account *Account, scope string, resetAt ti if reason = strings.TrimSpace(reason); reason != "" { payload["reason"] = reason } + + oldLimits, _ := account.Extra[modelRateLimitsKey].(map[string]any) + limits := make(map[string]any, len(oldLimits)+1) + maps.Copy(limits, oldLimits) limits[scope] = payload + + extra := make(map[string]any, len(account.Extra)+1) + maps.Copy(extra, account.Extra) + extra[modelRateLimitsKey] = limits + account.Extra = extra } diff --git a/backend/internal/service/model_rate_limit_snapshot_cow_test.go b/backend/internal/service/model_rate_limit_snapshot_cow_test.go new file mode 100644 index 000000000000..a73fcff5fe9a --- /dev/null +++ b/backend/internal/service/model_rate_limit_snapshot_cow_test.go @@ -0,0 +1,209 @@ +//go:build unit + +package service + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestSetAccountModelRateLimitSnapshot_CopyOnWrite(t *testing.T) { + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + existingLimit := map[string]any{ + "rate_limited_at": now.Format(time.RFC3339), + "rate_limit_reset_at": now.Add(time.Hour).Format(time.RFC3339), + } + limits := map[string]any{"claude-opus-5": existingLimit} + extra := map[string]any{ + modelRateLimitsKey: limits, + "session_window_utilization": 0.5, + } + account := &Account{ID: 1, Platform: PlatformAnthropic, Extra: extra} + before := *account + + setAccountModelRateLimitSnapshot(account, anthropicFableRateLimitKey, now.Add(2*time.Hour), " threshold ", now) + + // 调用前的 Extra 与嵌套 limits 原样保留。 + require.Len(t, extra, 2) + require.Equal(t, 0.5, extra["session_window_utilization"]) + require.Len(t, limits, 1) + require.Equal(t, existingLimit, limits["claude-opus-5"]) + require.NotContains(t, limits, anthropicFableRateLimitKey) + require.False(t, before.isModelRateLimitedWithContext(ctx, "claude-fable-5")) + + // 新 Extra 同时带有旧条目与新限流，本请求后续读取结果与原地写入一致。 + require.Equal(t, 0.5, account.Extra["session_window_utilization"]) + newLimits, ok := account.Extra[modelRateLimitsKey].(map[string]any) + require.True(t, ok) + require.Len(t, newLimits, 2) + require.Equal(t, existingLimit, newLimits["claude-opus-5"]) + require.Equal(t, map[string]any{ + "rate_limited_at": now.Format(time.RFC3339), + "rate_limit_reset_at": now.Add(2 * time.Hour).Format(time.RFC3339), + "reason": "threshold", + }, newLimits[anthropicFableRateLimitKey]) + require.Equal(t, "threshold", account.modelRateLimitReason(anthropicFableRateLimitKey)) + require.True(t, account.isModelRateLimitedWithContext(ctx, "claude-fable-5")) + require.True(t, account.isModelRateLimitedWithContext(ctx, "claude-opus-5")) + require.False(t, account.isModelRateLimitedWithContext(ctx, "claude-sonnet-5")) +} + +func TestSetAccountModelRateLimitSnapshot_InitializesMissingOrInvalidLimits(t *testing.T) { + ctx := context.Background() + now := time.Now().UTC() + + empty := &Account{Platform: PlatformAnthropic} + setAccountModelRateLimitSnapshot(empty, "claude-opus-5", now.Add(time.Hour), "", now) + require.True(t, empty.isModelRateLimitedWithContext(ctx, "claude-opus-5")) + require.Empty(t, empty.modelRateLimitReason("claude-opus-5")) + + invalidExtra := map[string]any{modelRateLimitsKey: "corrupted", "keep": true} + invalid := &Account{Platform: PlatformAnthropic, Extra: invalidExtra} + setAccountModelRateLimitSnapshot(invalid, "claude-opus-5", now.Add(time.Hour), "r", now) + require.True(t, invalid.isModelRateLimitedWithContext(ctx, "claude-opus-5")) + require.Equal(t, true, invalid.Extra["keep"]) + require.Equal(t, "corrupted", invalidExtra[modelRateLimitsKey]) + + unchanged := &Account{Platform: PlatformAnthropic, Extra: map[string]any{"k": 1}} + setAccountModelRateLimitSnapshot(unchanged, " ", now.Add(time.Hour), "r", now) + require.Equal(t, map[string]any{"k": 1}, unchanged.Extra) +} + +// 同一母本的浅拷贝并发写各自的限流快照：-race 下不得出现竞争，母本与其它副本都看不到这些写入。 +func TestSetAccountModelRateLimitSnapshot_ConcurrentCopiesOfSharedMaster(t *testing.T) { + ctx := context.Background() + now := time.Now().UTC() + master := &Account{ + ID: 7, + Platform: PlatformAnthropic, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + "claude-opus-5": map[string]any{"rate_limit_reset_at": now.Add(time.Hour).Format(time.RFC3339)}, + }, + "session_window_utilization": 0.2, + }, + } + master.PrepareSharedSnapshot() + + const workers = 16 + var wg sync.WaitGroup + for worker := 0; worker < workers; worker++ { + wg.Add(1) + go func(worker int) { + defer wg.Done() + own := fmt.Sprintf("model-%d", worker) + other := fmt.Sprintf("model-%d", (worker+1)%workers) + for i := 0; i < 50; i++ { + replica := *master + setAccountModelRateLimitSnapshot(&replica, own, now.Add(time.Hour), "r", now) + if !replica.isModelRateLimitedWithContext(ctx, own) { + t.Errorf("worker %d: own limit missing", worker) + return + } + if !replica.isModelRateLimitedWithContext(ctx, "claude-opus-5") { + t.Errorf("worker %d: inherited limit missing", worker) + return + } + if replica.isModelRateLimitedWithContext(ctx, other) { + t.Errorf("worker %d: saw another replica's limit", worker) + return + } + probe := *master + if probe.isModelRateLimitedWithContext(ctx, own) { + t.Errorf("worker %d: master was mutated", worker) + return + } + } + }(worker) + } + wg.Wait() + + limits, ok := master.Extra[modelRateLimitsKey].(map[string]any) + require.True(t, ok) + require.Len(t, limits, 1) + require.Len(t, master.Extra, 2) +} + +type sharedSnapshotFableRepoStub struct { + AccountRepository + modelCalls atomic.Int64 +} + +func (r *sharedSnapshotFableRepoStub) SetModelRateLimit(_ context.Context, _ int64, _ string, _ time.Time, _ ...string) error { + r.modelCalls.Add(1) + return nil +} + +// Fable 阈值在选号阶段写入候选对象：同一母本的其它副本（包括并发读取的副本）与母本本身都不得出现该限流。 +func TestApplyAccountSchedulingThreshold_FableLimitStaysOnOwnSnapshotCopy(t *testing.T) { + accountSchedulingThresholdsSF.Forget(SettingKeyAccountSchedulingThresholds) + accountSchedulingThresholdsCache.Store(&cachedAccountSchedulingThresholds{}) + t.Cleanup(func() { + accountSchedulingThresholdsSF.Forget(SettingKeyAccountSchedulingThresholds) + accountSchedulingThresholdsCache.Store(&cachedAccountSchedulingThresholds{}) + }) + + settingsRepo := newMockSettingRepo() + settingsRepo.data[SettingKeyAccountSchedulingThresholds] = `{"anthropic":100}` + repo := &sharedSnapshotFableRepoStub{} + limiter := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + limiter.SetSettingService(NewSettingService(settingsRepo, &config.Config{})) + + now := time.Now().UTC() + master := &Account{ + ID: 1006, + Platform: PlatformAnthropic, + Status: StatusActive, + Schedulable: true, + Credentials: map[string]any{"account_scheduling_threshold": 60}, + Extra: map[string]any{ + "passive_usage_7d_utilization": 0.40, + "passive_usage_7d_reset": float64(now.Add(3 * 24 * time.Hour).Unix()), + "passive_usage_7d_oi_utilization": 0.61, + "passive_usage_7d_oi_reset": float64(now.Add(4 * 24 * time.Hour).Unix()), + }, + } + master.PrepareSharedSnapshot() + + const pairs = 8 + var wg sync.WaitGroup + for i := 0; i < pairs; i++ { + wg.Add(2) + go func() { + defer wg.Done() + candidate := *master + if limiter.ApplyAccountSchedulingThreshold(context.Background(), &candidate) { + t.Error("the Fable-only window must not pause the whole account") + } + if candidate.IsSchedulableForModel("claude-fable-5") { + t.Error("the candidate that crossed the threshold must stop scheduling Fable") + } + if !candidate.IsSchedulableForModel("claude-opus-5") { + t.Error("non-Fable models must stay schedulable") + } + }() + go func() { + defer wg.Done() + for j := 0; j < 20; j++ { + reader := *master + if !reader.IsSchedulableForModel("claude-fable-5") { + t.Error("a sibling copy must not see the Fable limit") + return + } + } + }() + } + wg.Wait() + + require.Equal(t, int64(pairs), repo.modelCalls.Load()) + require.NotContains(t, master.Extra, modelRateLimitsKey) + require.True(t, master.IsSchedulableForModel("claude-fable-5")) +} diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index c3040ebf9ca2..77a3bdf98dd4 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -456,7 +456,7 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont if isOpenAIRequestScopedSafetyRejection(responseBody) { return false } - if isOpenAIRequestBudgetRejection(account, statusCode, responseBody) || isOpenAIReportedUpstreamFailure(statusCode, responseBody) { + if isOpenAIRequestBudgetRejection(account, statusCode, responseBody) { return false } if account != nil && account.Platform == PlatformGrok && isGrokContentPolicyRejection(statusCode, responseBody) { diff --git a/backend/internal/service/openai_capacity_shed_test.go b/backend/internal/service/openai_capacity_shed_test.go index 7d2f8c8dd28a..6737e42cbdf9 100644 --- a/backend/internal/service/openai_capacity_shed_test.go +++ b/backend/internal/service/openai_capacity_shed_test.go @@ -3,6 +3,7 @@ package service import ( "context" "errors" + "fmt" "io" "net/http" "net/http/httptest" @@ -14,6 +15,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/openai" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" ) // --- mock: 只记录临时不可调度写入，其余方法不应被调用 --- @@ -108,6 +110,33 @@ func TestOpenAIHTTPCapacityShedIsRequestScopedForOAuthAccounts(t *testing.T) { require.Zero(t, repo.tempUnschedCalls) } +func TestOpenAICapacityShedHonorsRetryAfter(t *testing.T) { + for _, test := range []struct { + name string + headers http.Header + body string + want time.Duration + }{ + {name: "http seconds", headers: http.Header{"Retry-After": []string{"6"}}, want: 6 * time.Second}, + {name: "http milliseconds", headers: http.Header{"Retry-After-Ms": []string{"6000"}}, want: 6 * time.Second}, + {name: "body seconds", body: `{"error":{"headers":{"retry-after":"7"},"message":"Our servers are currently overloaded. Please try again later."}}`, want: 7 * time.Second}, + {name: "body milliseconds", body: `{"response":{"error":{"headers":{"retry-after-ms":"7000"},"message":"Our servers are currently overloaded. Please try again later."}}}`, want: 7 * time.Second}, + {name: "bounded", headers: http.Header{"Retry-After": []string{"60"}}, want: openAIRequestScopedRetryMaxDelay}, + } { + t.Run(test.name, func(t *testing.T) { + failoverErr := newOpenAIUpstreamFailoverError( + http.StatusServiceUnavailable, + test.headers, + []byte(test.body), + "Our servers are currently overloaded. Please try again later.", + false, + ) + require.True(t, failoverErr.RequestScopedTransient) + require.Equal(t, test.want, failoverErr.SameAccountRetryDelay) + }) + } +} + // 上游降载的真实序列是「event: error → event: response.failed」。error 帧不算 // 客户端输出：若把它当首输出 flush，clientOutputStarted 被固化，随后的 failed // 事件就进不了 pre-output failover 分支，只能把致命错误原样转发给客户端。 @@ -124,6 +153,8 @@ func TestOpenAIStreamErrorFrameDoesNotStartClientOutput(t *testing.T) { {`{"type":"error","error":{"type":"invalid_request_error","code":"content_policy_violation","message":"blocked"}}`, "error", true}, {`{"type":"response.failed","response":{"error":{"code":"server_is_overloaded"}}}`, "response.failed", false}, {`{"type":"response.created","response":{"id":"resp_1"}}`, "response.created", false}, + {`{"type":"response.queued","response":{"id":"resp_1","status":"queued"}}`, "response.queued", false}, + {`{"type":"response.metadata","response_id":"resp_1","metadata":{"openai_chatgpt_moderation_metadata":{"prompt":{}}}}`, "response.metadata", false}, {`{"type":"response.in_progress","response":{"id":"resp_1"}}`, "response.in_progress", false}, {`{"type":"keepalive"}`, "keepalive", false}, {`{"type":"unknown_event"}`, "unknown_event", true}, @@ -147,6 +178,12 @@ func TestOpenAIStreamMetadataAndKeepaliveBeforeOverloadFailOver(t *testing.T) { "event: response.created", `data: {"type":"response.created","response":{"id":"resp_1","metadata":{"padding":"` + largeMetadata + `"}}}`, "", + "event: response.queued", + `data: {"type":"response.queued","response":{"id":"resp_1","status":"queued"}}`, + "", + "event: response.metadata", + `data: {"type":"response.metadata","response_id":"resp_1","metadata":{"openai_chatgpt_moderation_metadata":{"prompt":{}}}}`, + "", "event: response.output_item.added", `data: {"type":"response.output_item.added","item":{"type":"reasoning","summary":[]}}`, "", @@ -210,6 +247,69 @@ func TestOpenAIStreamMetadataAndKeepaliveBeforeOverloadFailOver(t *testing.T) { } } +func TestOpenAIStreamEmptyContentBeforeFailureCanFailOver(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, payload := range []string{ + `{"type":"response.output_text.delta","delta":""}`, + `{"type":"response.reasoning_summary_text.delta","delta":""}`, + `{"type":"response.reasoning_text.delta","delta":""}`, + `{"type":"response.function_call_arguments.delta","delta":""}`, + `{"type":"response.custom_tool_call_input.delta","delta":""}`, + `{"type":"response.output_text.done","text":""}`, + `{"type":"response.reasoning_summary_text.done","text":""}`, + `{"type":"response.reasoning_text.done","text":""}`, + `{"type":"response.refusal.delta","delta":""}`, + `{"type":"response.refusal.done","refusal":""}`, + `{"type":"response.content_part.done","part":{"type":"output_text","text":""}}`, + `{"type":"response.reasoning_summary_part.done","part":{"type":"summary_text","text":""}}`, + } { + for _, passthrough := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/passthrough=%t", payload, passthrough), func(t *testing.T) { + svc := &OpenAIGatewayService{cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}} + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader("data: " + payload + "\n\n" + + `data: {"type":"response.failed","response":{"status":"failed","error":{"code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}}}` + "\n\n")), + } + account := &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth} + var err error + if passthrough { + _, err = svc.handleStreamingResponsePassthrough(c.Request.Context(), resp, c, account, time.Now(), "model", "model") + } else { + _, err = svc.handleStreamingResponse(c.Request.Context(), resp, c, account, time.Now(), "model", "model") + } + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.True(t, failoverErr.RequestScopedTransient) + require.Empty(t, rec.Body.String()) + }) + } + } +} + +func TestOpenAIStreamReplayBoundaryPreservesContentAndToolCompletion(t *testing.T) { + for _, payload := range []string{ + `{"type":"response.output_text.delta","delta":" "}`, + `{"type":"response.function_call_arguments.delta","delta":"{}"}`, + `{"type":"response.custom_tool_call_input.delta","delta":"run"}`, + `{"type":"response.function_call_arguments.done","arguments":""}`, + `{"type":"response.custom_tool_call_input.done","input":""}`, + `{"type":"response.output_item.done","item":{"type":"function_call","arguments":""}}`, + `{"type":"response.output_item.done","item":{"type":"reasoning","encrypted_content":"ciphertext"}}`, + `{"type":"response.output_text.delta","delta":null}`, + `{"type":"response.output_text.delta"}`, + `{"type":"response.output_text.delta","delta":{}}`, + `{"type":"response.new_kind.delta","delta":""}`, + } { + eventType := gjson.Get(payload, "type").String() + require.True(t, openAIStreamDataStartsClientOutput(payload, eventType), payload) + } +} + func TestOpenAIStreamKeepaliveOnlyFailureDoesNotRecordFirstToken(t *testing.T) { gin.SetMode(gin.TestMode) for _, passthrough := range []bool{false, true} { diff --git a/backend/internal/service/openai_codex_client_identity.go b/backend/internal/service/openai_codex_client_identity.go index 85540c8c7cde..fec73be34417 100644 --- a/backend/internal/service/openai_codex_client_identity.go +++ b/backend/internal/service/openai_codex_client_identity.go @@ -15,7 +15,7 @@ import ( // CodexClientIdentityExtraKey 持久化在 accounts.extra 中的每账号 Codex 客户端身份。 // OpenAI OAuth（及 setup_token）账号一账号一身份：由系统种子确定性派生后固定， // 所有出站请求（HTTP / 透传 / WS 握手 / 探针 / token 刷新）共用，使上游看到的 -// 是同一台设备上运行的同一个真实 Codex TUI，而不是多个用户各自的客户端。 +// 是同一台固定设备；出站客户端标识由统一身份解析器归一。 const CodexClientIdentityExtraKey = "codex_client_identity" const codexClientIdentitySchemaVersion = 1 @@ -105,7 +105,7 @@ func codexSandboxForUserAgent(userAgent string) string { } } -// UserAgent 拼出真实 Codex TUI 的 User-Agent： +// UserAgent 保留设备原有的 Codex TUI User-Agent 形态，出站由官方规则归一为 CLI： // codex-tui/ ( ; ) (codex-tui; )。 // 首段版本、尾部括号组版本与 version 头必须同源（codex-rs 三处都取同一个 // CARGO_PKG_VERSION）。 diff --git a/backend/internal/service/openai_codex_client_identity_test.go b/backend/internal/service/openai_codex_client_identity_test.go index f074aa4f14d8..9148793cf61e 100644 --- a/backend/internal/service/openai_codex_client_identity_test.go +++ b/backend/internal/service/openai_codex_client_identity_test.go @@ -53,7 +53,7 @@ func TestRefreshAccountTokenUsesAccountCodexIdentity(t *testing.T) { require.Equal(t, "with-identity", info.AccessToken) expected := resolveCodexOutboundIdentityForAccount(account, "") require.Equal(t, expected.userAgent, stub.userAgent) - require.Equal(t, "codex-tui", stub.originator) + require.Equal(t, openai.CodexCLIOriginator, stub.originator) } func TestCodexIdentitySnapshotIsSecretFreeAndConsistent(t *testing.T) { @@ -74,8 +74,9 @@ func TestCodexIdentitySnapshotIsSecretFreeAndConsistent(t *testing.T) { SetCodexForceCLIEnabled(true) forced := resolveCodexIdentitySnapshot(account, account, codexAccountIdentityOverrideUA(account)) require.Equal(t, "account", forced.IdentitySource) - require.Equal(t, deriveCodexClientIdentity(seed).UserAgent(forced.Version), forced.UserAgent, - "ForceCodexCLI drops the custom UA while retaining the account's derived TUI profile") + _, expectedUA, _ := openai.NormalizeCodexClientIdentityToCLI(codexTUIOriginator, deriveCodexClientIdentity(seed).UserAgent(forced.Version)) + require.Equal(t, expectedUA, forced.UserAgent, + "ForceCodexCLI drops the custom UA while retaining the account's derived device profile") } func TestCodexIdentitySnapshotReportsTheSelectedFingerprintSource(t *testing.T) { @@ -107,9 +108,10 @@ func TestCodexAccountIdentityUsesBuiltInVersionWhenCanonicalVersionIsInvalid(t * Extra: map[string]any{codexFingerprintSeedExtraKey: "1c0a3d9e-58b2-4f8c-a2d1-7f3b9e6c4a55"}} profile, ok := account.CodexClientIdentity() require.True(t, ok) + originator, userAgent, _ := openai.NormalizeCodexClientIdentityToCLI(codexTUIOriginator, profile.UserAgent(codexCLIVersion)) expected := codexOutboundIdentity{ - userAgent: profile.UserAgent(codexCLIVersion), - originator: codexTUIOriginator, + userAgent: userAgent, + originator: originator, version: codexCLIVersion, } @@ -172,9 +174,8 @@ func TestCodexIdentityBackfillWritesOnlyMissingIdentities(t *testing.T) { } } -// 账号级身份：由种子确定性派生，UA 三处版本同源，形态与 codex-rs 的 -// `codex-tui/ ( ; ) (codex-tui; )` 一致。 -func TestEnforceCodexIdentityHeadersForAccountUsesDerivedTUIIdentity(t *testing.T) { +// 账号级设备由种子确定性派生；出站采用官方 CLI 标识，保留设备及有效版本。 +func TestEnforceCodexIdentityHeadersForAccountUsesDerivedCLIIdentity(t *testing.T) { seed := "6f2c1c7e-6e2d-4a4b-9f0b-2f5f0c8a1d33" account := &Account{ID: 5, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{"chatgpt_account_id": "acct-5"}, @@ -192,11 +193,12 @@ func TestEnforceCodexIdentityHeadersForAccountUsesDerivedTUIIdentity(t *testing. enforceCodexIdentityHeadersForAccount(h, account, "") version := resolveCodexOutboundIdentity("").version - uaPattern := regexp.MustCompile(`^codex-tui/` + regexp.QuoteMeta(version) + - ` \((Mac OS|Windows|Ubuntu) [0-9]+\.[0-9]+\.[0-9]+; (arm64|x86_64)\) \S+ \(codex-tui; ` + regexp.QuoteMeta(version) + `\)$`) + uaPattern := regexp.MustCompile(`^codex_cli_rs/` + regexp.QuoteMeta(version) + + ` \((Mac OS|Windows|Ubuntu) [0-9]+\.[0-9]+\.[0-9]+; (arm64|x86_64)\) \S+$`) require.Regexp(t, uaPattern, h.Get("user-agent")) - require.Equal(t, first.UserAgent(version), h.Get("user-agent")) - require.Equal(t, "codex-tui", h.Get("originator")) + require.Contains(t, h.Get("user-agent"), fmt.Sprintf("(%s %s; %s) %s", first.OSType, first.OSVersion, first.Arch, first.Terminal)) + require.NotContains(t, h.Get("user-agent"), "codex-tui") + require.Equal(t, openai.CodexCLIOriginator, h.Get("originator")) require.Equal(t, version, h.Get("version")) require.Contains(t, []string{"seatbelt", "windows_sandbox", "seccomp"}, first.Sandbox) } diff --git a/backend/internal/service/openai_codex_identity.go b/backend/internal/service/openai_codex_identity.go index 2d612e10b3ef..1052710a700b 100644 --- a/backend/internal/service/openai_codex_identity.go +++ b/backend/internal/service/openai_codex_identity.go @@ -32,7 +32,7 @@ func NormalizeCodexClientVersion(version string) string { return version } -// buildCodexCLIUserAgent 按版本号拼出规范 Codex TUI User-Agent。 +// buildCodexCLIUserAgent 按版本号拼出规范 Codex CLI User-Agent。 // UA 形态只在 codexCLIUserAgentSuffix 一处定义，避免多处拼装漂移。 func buildCodexCLIUserAgent(version string) string { if version = NormalizeCodexClientVersion(version); version == "" { @@ -43,7 +43,7 @@ func buildCodexCLIUserAgent(version string) string { // codexIdentityEnforcement 控制 enforceCodexIdentityHeaders 是否强制统一出站身份， // 由 gateway.disable_codex_identity_enforcement 在服务构造时取反发布。 -// 默认开启：上游在容量紧张时按客户端身份分优先级降载，被降载的请求会拿到 +// 默认开启：上游会按客户端身份进行风控分层，命中风险身份的请求会拿到 // HTTP 200 + 流内 server_is_overloaded，本次请求即失败；强制统一出口可确保没有 // 请求带着第三方或陈旧身份出站。关闭后退回「仅按最终 UA 配对 originator」的收口语义。 var codexIdentityEnforcement = func() *atomic.Bool { @@ -154,11 +154,11 @@ type codexOutboundIdentity struct { } // resolveCodexOutboundIdentity 由候选 User-Agent 推导自洽的出站身份。 -// candidateUA 为空时使用规范 User-Agent；推导不出官方身份时整体回退为规范 TUI 身份。 +// candidateUA 为空时使用规范 User-Agent；推导不出官方身份时整体回退为规范 CLI 身份。 // // 候选 UA（面板 / 账号级的管理员显式配置）只贡献客户端名与 OS / 架构 / 终端指纹， // 其自带的版本段一律用当前生效版本重建：一条填写于某个历史版本的 UA 否则会把出站身份 -// 永久钉死在陈旧版本上，绕过版本自动同步，落回上游优先降载的那一侧。 +// 永久钉死在陈旧版本上，绕过版本自动同步，落回上游身份风控更严格的一侧。 // 需要固定版本请填「Codex 客户端版本号」并关闭自动同步。 func resolveCodexOutboundIdentity(candidateUA string, policies ...*codexFingerprintPolicy) codexOutboundIdentity { policy := selectCodexFingerprintPolicy(policies) @@ -176,9 +176,12 @@ func resolveCodexOutboundIdentity(candidateUA string, policies ...*codexFingerpr originator, pairedUA, ok := openai.PairCodexClientIdentity(ua) if !ok { if originator, pairedUA, ok = openai.PairCodexClientIdentity(canonical); !ok { - originator, pairedUA = openai.CodexDefaultOriginator, codexCLIUserAgent + originator, pairedUA = openai.CodexCLIOriginator, codexCLIUserAgent } } + // 旧配置或账号级 UA 可能仍保存 codex-tui。该身份是 sub2api 默认出口曾经 + // 使用的风险指纹，必须与 UA 首段一起归一为 CLI，不能只替换 originator。 + originator, pairedUA, _ = openai.NormalizeCodexClientIdentityToCLI(originator, pairedUA) // 生效版本只有一个来源：规范身份（面板版本号 → 自动同步值 → 内置常量，见 // SettingService.GetOpenAICodexClientVersion）。UA 与 version 头由此同源派生。 version := codexClientVersionFromUA(canonical) @@ -243,8 +246,8 @@ func enforceCodexIdentityHeaders(h http.Header) { } // enforceCodexIdentityHeadersWithUA 强制统一 OAuth 出站身份：User-Agent / originator / version -// 一律改写为网关的规范身份，客户端自报身份不参与构造。上游在容量紧张时按客户端身份分优先级 -// 降载，被降载的请求会拿到 HTTP 200 + 流内 server_is_overloaded；统一出口可确保没有请求带着 +// 一律改写为网关的规范身份，客户端自报身份不参与构造。上游会按客户端身份进行风控分层， +// 命中风险身份的请求会拿到 HTTP 200 + 流内 server_is_overloaded；统一出口可确保没有请求带着 // 第三方或陈旧身份出站，也天然满足 originator 与 UA 首段配套的上游校验（issue #3901）。 // // overrideUA 是账号级自定义 User-Agent：管理员的显式配置仍然生效，但只贡献客户端名与 @@ -272,8 +275,8 @@ func enforceCodexIdentityHeadersWithUA(h http.Header, overrideUA string, policie // resolveCodexOutboundIdentityForAccount 返回某个 OAuth 账号的出站身份三元组。 // 优先级：管理员显式配置的账号级 User-Agent（只贡献客户端名与 OS / 架构 / 终端指纹， -// 版本段仍由生效版本重建）> 账号持久化 / 种子派生的 Codex TUI 身份 > 全局规范身份。 -// 版本号三处同源：UA 首段、UA 尾部括号组与 version 头都取当前生效的官方版本。 +// 版本段仍由生效版本重建）> 账号持久化 / 种子派生的设备身份 > 全局规范身份。 +// 账号设备沿用官方 CLI 归一化，版本号与 version 头都取当前生效的官方版本。 func resolveCodexOutboundIdentityForAccount(account *Account, overrideUA string, policies ...*codexFingerprintPolicy) codexOutboundIdentity { policy := selectCodexFingerprintPolicy(policies) canonical := resolveCodexOutboundIdentity(overrideUA, policy) @@ -287,16 +290,17 @@ func resolveCodexOutboundIdentityForAccount(account *Account, overrideUA string, if !ok { return canonical } + originator, userAgent, _ := openai.NormalizeCodexClientIdentityToCLI(codexTUIOriginator, identity.UserAgent(canonical.version)) return codexOutboundIdentity{ - userAgent: identity.UserAgent(canonical.version), - originator: codexTUIOriginator, + userAgent: userAgent, + originator: originator, version: canonical.version, } } // enforceCodexIdentityHeadersForAccount 与 enforceCodexIdentityHeadersWithUA 语义相同， // 但强制统一时使用账号级身份而不是全局规范身份，使同一账号的所有出站请求 -// 表现为同一台机器上的同一个 Codex TUI。 +// 使用同一台固定设备，并采用官方 CLI 标识。 func enforceCodexIdentityHeadersForAccount(h http.Header, account *Account, overrideUA string, policies ...*codexFingerprintPolicy) { if h == nil || h.Get("originator") == "" { return diff --git a/backend/internal/service/openai_codex_identity_snapshot.go b/backend/internal/service/openai_codex_identity_snapshot.go index 602bda643263..fdf1ab972499 100644 --- a/backend/internal/service/openai_codex_identity_snapshot.go +++ b/backend/internal/service/openai_codex_identity_snapshot.go @@ -123,7 +123,7 @@ func resolveCodexIdentitySnapshot(routed, source *Account, overrideUA string, po snapshot.IdentitySource = "override_ua" } } else if source != nil { - if accountIdentity, ok := source.CodexClientIdentity(); ok && accountIdentity.UserAgent(identity.version) == identity.userAgent { + if _, ok := source.CodexClientIdentity(); ok { snapshot.IdentitySource = "account" } } diff --git a/backend/internal/service/openai_codex_identity_test.go b/backend/internal/service/openai_codex_identity_test.go index 374853d7a217..a078a6b19516 100644 --- a/backend/internal/service/openai_codex_identity_test.go +++ b/backend/internal/service/openai_codex_identity_test.go @@ -174,8 +174,8 @@ func TestEnforceCodexIdentityHeadersWithAccountOverrideUA(t *testing.T) { enforceCodexIdentityHeadersWithUA(h, "codex-tui/0.125.0 (Mac OS X 14.0; arm64) iTerm") - require.Equal(t, "codex-tui", h.Get("originator")) - require.Equal(t, "codex-tui/0.200.1 (Mac OS X 14.0; arm64) iTerm", h.Get("user-agent")) + require.Equal(t, openai.CodexCLIOriginator, h.Get("originator")) + require.Equal(t, "codex_cli_rs/0.200.1 (Mac OS X 14.0; arm64) iTerm", h.Get("user-agent")) require.Equal(t, "0.200.1", h.Get("version")) }) } diff --git a/backend/internal/service/openai_codex_version_sync_service_test.go b/backend/internal/service/openai_codex_version_sync_service_test.go index 7d40593e53b5..cdd0be6ddec0 100644 --- a/backend/internal/service/openai_codex_version_sync_service_test.go +++ b/backend/internal/service/openai_codex_version_sync_service_test.go @@ -7,6 +7,7 @@ import ( "testing" "time" + "github.com/Wei-Shaw/sub2api/internal/pkg/openai" "github.com/stretchr/testify/require" ) @@ -327,14 +328,14 @@ func TestGetOpenAICodexClientVersionFallsBackOnError(t *testing.T) { require.Equal(t, codexCLIVersion, svc.GetOpenAICodexClientVersion(context.Background())) } -// 规范 UA：面板未填完整 UA 时按当前生效版本号拼出标准 TUI 形态。 +// 规范 UA：面板未填完整 UA 时按当前生效版本号拼出标准 CLI 形态。 func TestGetOpenAICodexCanonicalUserAgentBuildsFromVersion(t *testing.T) { svc := NewSettingService(&codexVersionSettingRepoStub{values: map[string]string{ SettingKeyOpenAICodexClientVersionSynced: "0.200.1", }}, nil) require.Equal(t, - "codex-tui/0.200.1"+codexCLIUserAgentSuffix, + "codex_cli_rs/0.200.1"+codexCLIUserAgentSuffix, svc.GetOpenAICodexCanonicalUserAgent(context.Background()), ) } @@ -517,16 +518,19 @@ func TestGetOpenAICodexCanonicalUserAgentOutboundIdentity(t *testing.T) { SettingKeyOpenAICodexClientVersion: version.manual, SettingKeyOpenAICodexClientVersionSynced: version.synced, }}, nil) - wantUA := buildCodexCLIUserAgent(version.want) - wantOriginator := "codex-tui" + wantCanonicalUA := buildCodexCLIUserAgent(version.want) + wantUA := wantCanonicalUA + wantOriginator := openai.CodexDefaultOriginator if ua.client != "" { wantUA = ua.client + "/" + version.want + ua.suffix + wantCanonicalUA = wantUA wantOriginator = ua.client } + wantOriginator, wantUA, _ = openai.NormalizeCodexClientIdentityToCLI(wantOriginator, wantUA) SetCodexCanonicalUserAgentResolver(func() string { return svc.GetOpenAICodexCanonicalUserAgent(context.Background()) }) - require.Equal(t, wantUA, svc.GetOpenAICodexCanonicalUserAgent(context.Background())) + require.Equal(t, wantCanonicalUA, svc.GetOpenAICodexCanonicalUserAgent(context.Background())) identity := resolveCodexOutboundIdentity("") require.Equal(t, wantUA, identity.userAgent) require.Equal(t, wantOriginator, identity.originator) diff --git a/backend/internal/service/openai_gateway_anthropic_native_pump.go b/backend/internal/service/openai_gateway_anthropic_native_pump.go index 8a1967770568..3b679bc4e8cd 100644 --- a/backend/internal/service/openai_gateway_anthropic_native_pump.go +++ b/backend/internal/service/openai_gateway_anthropic_native_pump.go @@ -1,8 +1,10 @@ package service -// 国产供应商 Anthropic 协议转换路径的上游 SSE 行泵。 +// Anthropic 协议转换路径的上游 SSE 行泵，用于 OpenAI 侧国产供应商的 +// CC×anthropic / Responses×anthropic 转换链，以及 Anthropic 账号的 +// chat/responses 兼容入口（anthropicCompatStream）。 // -// 这两条转换链（CC×anthropic / Responses×anthropic）的上游 ctx 是 +// 这些转换链的上游 ctx 是 // WithoutCancel（detachStreamUpstreamContext）、http.Client 无整体 Timeout， // 客户端断开后的排水阶段若上游挂住 SSE（不发数据也不断连）， // scanner.Scan() 将永久阻塞：goroutine 钉死、resp.Body 不归还、连接池位 @@ -18,11 +20,17 @@ import ( "errors" "io" "time" + + "github.com/Wei-Shaw/sub2api/internal/config" ) // errAnthropicNativeStreamIdle 表示上游流读间隔超时（见上方文件注释）。 var errAnthropicNativeStreamIdle = errors.New("stream data interval timeout") +// errAnthropicNativeClientGone 表示等待上游数据期间客户端已离开；泵状态不变， +// 调用方可继续读取。 +var errAnthropicNativeClientGone = errors.New("client gone") + // anthropicNativeLineEvent 是行泵交付的单次读取结果：line 为一行 SSE 文本， // err 为 scanner 读错误（流自然结束时 next 返回 io.EOF，不经过本字段）。 type anthropicNativeLineEvent struct { @@ -31,7 +39,7 @@ type anthropicNativeLineEvent struct { } // anthropicNativeLinePump 以独立 goroutine 泵送 scanner 的行，并对逐行到达 -// 间隔施加 interval 上限（<=0 表示禁用，保持无界读的旧行为）。 +// 间隔施加 interval 上限（<=0 表示禁用，可经 enableInterval 启用）。 type anthropicNativeLinePump struct { events chan anthropicNativeLineEvent done chan struct{} @@ -69,23 +77,51 @@ func newAnthropicNativeLinePump(scanner *bufio.Scanner, interval time.Duration) } // next 阻塞返回下一行。返回 io.EOF 表示上游正常收流；errAnthropicNativeStreamIdle -// 表示 interval 内无任何数据到达（计时从收到上一行时起算，事件处理耗时不算入， -// 与 readOpenAICompatBufferedTerminal 的 resetTimeout 语义一致）。 +// 表示自交付上一行起已过 interval 且没有可交付的行；计时器到期时若已有缓冲的行， +// 优先交付该行，调用方处理事件期间上游送达的数据不会被判为上游空闲。 func (p *anthropicNativeLinePump) next() (string, error) { + return p.nextOrClientGone(nil) +} + +// nextOrClientGone 与 next 相同，另在等待期间 clientGone 关闭时返回 +// errAnthropicNativeClientGone。clientGone 为 nil 时不监听。 +func (p *anthropicNativeLinePump) nextOrClientGone(clientGone <-chan struct{}) (string, error) { var timeoutCh <-chan time.Time if p.timer != nil { timeoutCh = p.timer.C } select { case ev, ok := <-p.events: - if !ok { - return "", io.EOF - } - p.resetTimer() - return ev.line, ev.err + return p.deliver(ev, ok) case <-timeoutCh: - return "", errAnthropicNativeStreamIdle + select { + case ev, ok := <-p.events: + return p.deliver(ev, ok) + default: + return "", errAnthropicNativeStreamIdle + } + case <-clientGone: + return "", errAnthropicNativeClientGone + } +} + +// deliver 交付一次读取结果；ok=false 表示泵 goroutine 已结束。 +func (p *anthropicNativeLinePump) deliver(ev anthropicNativeLineEvent, ok bool) (string, error) { + if !ok { + return "", io.EOF + } + p.resetTimer() + return ev.line, ev.err +} + +// enableInterval 在读间隔上限未启用时以 interval 启用，计时从调用时起算； +// 已启用时不变。只能在调用 next 的 goroutine 中调用。 +func (p *anthropicNativeLinePump) enableInterval(interval time.Duration) { + if p.timer != nil || interval <= 0 { + return } + p.interval = interval + p.timer = time.NewTimer(interval) } // resetTimer 在收到一行后重启间隔计时器。 @@ -116,11 +152,11 @@ func (p *anthropicNativeLinePump) stop() { } } -// anthropicNativeStreamInterval 返回本组转换路径适用的读间隔上限； +// anthropicNativeStreamInterval 返回 Anthropic 上游协议转换路径适用的读间隔上限； // gateway.stream_data_interval_timeout <= 0 时视为禁用。 -func (s *OpenAIGatewayService) anthropicNativeStreamInterval() time.Duration { - if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { - return time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second +func anthropicNativeStreamInterval(cfg *config.Config) time.Duration { + if cfg != nil && cfg.Gateway.StreamDataIntervalTimeout > 0 { + return time.Duration(cfg.Gateway.StreamDataIntervalTimeout) * time.Second } return 0 } diff --git a/backend/internal/service/openai_gateway_anthropic_native_pump_test.go b/backend/internal/service/openai_gateway_anthropic_native_pump_test.go index 736e8f7ae42f..103137a33b4c 100644 --- a/backend/internal/service/openai_gateway_anthropic_native_pump_test.go +++ b/backend/internal/service/openai_gateway_anthropic_native_pump_test.go @@ -130,6 +130,87 @@ func TestAnthropicNativeLinePump_DataResetsTimer(t *testing.T) { } } +func TestAnthropicNativeLinePump_SlowConsumerGetsBufferedLineInsteadOfIdle(t *testing.T) { + // 计时器与缓冲行同时就绪时 select 随机选择，多轮重复以覆盖两种选择。 + for round := 0; round < 64; round++ { + pump := &anthropicNativeLinePump{ + events: make(chan anthropicNativeLineEvent, 1), + done: make(chan struct{}), + timer: time.NewTimer(0), + interval: time.Hour, + } + pump.events <- anthropicNativeLineEvent{line: "data: {}"} + time.Sleep(time.Millisecond) + + line, err := pump.next() + pump.stop() + if err != nil || line != "data: {}" { + t.Fatalf("round %d: expected buffered line, got %q err=%v", round, line, err) + } + } +} + +func TestAnthropicNativeLinePump_ClientGoneKeepsPumpReadable(t *testing.T) { + pr, pw := io.Pipe() + scanner := bufio.NewScanner(pr) + pump := newAnthropicNativeLinePump(scanner, 0) + defer pump.stop() + defer func() { _ = pr.Close() }() + + clientGone := make(chan struct{}) + close(clientGone) + if _, err := pump.nextOrClientGone(clientGone); err != errAnthropicNativeClientGone { + t.Fatalf("expected client gone, got %v", err) + } + + go func() { + _, _ = pw.Write([]byte("event: ping\n")) + _ = pw.Close() + }() + line, err := pump.nextOrClientGone(nil) + if err != nil || line != "event: ping" { + t.Fatalf("expected the pending line after client gone, got %q err=%v", line, err) + } + if _, err := pump.next(); err != io.EOF { + t.Fatalf("expected EOF, got %v", err) + } +} + +func TestAnthropicNativeLinePump_EnableIntervalBoundsDisabledPump(t *testing.T) { + pr, _ := io.Pipe() + scanner := bufio.NewScanner(pr) + defer func() { _ = pr.Close() }() + + pump := newAnthropicNativeLinePump(scanner, 0) + defer pump.stop() + pump.enableInterval(50 * time.Millisecond) + + start := time.Now() + if _, err := pump.next(); err != errAnthropicNativeStreamIdle { + t.Fatalf("expected interval timeout after enableInterval, got %v", err) + } + if elapsed := time.Since(start); elapsed > 5*time.Second { + t.Fatalf("timeout not respected: %v", elapsed) + } +} + +func TestAnthropicNativeLinePump_EnableIntervalKeepsConfiguredInterval(t *testing.T) { + pr, _ := io.Pipe() + scanner := bufio.NewScanner(pr) + defer func() { _ = pr.Close() }() + + pump := newAnthropicNativeLinePump(scanner, 50*time.Millisecond) + defer pump.stop() + pump.enableInterval(time.Hour) + + if pump.interval != 50*time.Millisecond { + t.Fatalf("configured interval must not change, got %v", pump.interval) + } + if _, err := pump.next(); err != errAnthropicNativeStreamIdle { + t.Fatalf("expected interval timeout, got %v", err) + } +} + func TestCCStreamingFromNativeAnthropic_HangTimesOut(t *testing.T) { gin.SetMode(gin.TestMode) svc := newNativeAnthropicHangTestService(1) diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index ff4cdff21801..e31619d1b5ad 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -1077,19 +1077,6 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse( } chunks := apicompat.ResponsesEventToChatChunks(&event, state) - if state.ProtocolError != "" { - streamNonFailoverErr = errors.New(state.ProtocolError) - if !clientDisconnected { - if !c.Writer.Written() { - writeChatCompletionsError(c, http.StatusBadGateway, "upstream_error", state.ProtocolError) - } else { - _, _ = fmt.Fprint(c.Writer, buildChatStreamErrorSSE("upstream_protocol_error", state.ProtocolError)+"data: [DONE]\n\n") - c.Writer.Flush() - } - clientDisconnected = true - } - return true - } if !clientDisconnected { for _, chunk := range chunks { semanticChunk := openAIChatChunksHaveSemanticOutput([]apicompat.ChatCompletionsChunk{chunk}) diff --git a/backend/internal/service/openai_gateway_chat_completions_anthropic_native.go b/backend/internal/service/openai_gateway_chat_completions_anthropic_native.go index afea4622b0db..122a65318e08 100644 --- a/backend/internal/service/openai_gateway_chat_completions_anthropic_native.go +++ b/backend/internal/service/openai_gateway_chat_completions_anthropic_native.go @@ -169,7 +169,7 @@ func (s *OpenAIGatewayService) handleCCBufferedFromNativeAnthropic( var usage ClaudeUsage // 读间隔上限：上游挂住 SSE 时中止组装（缓冲路径尚未提交响应头，可回 502）。 - streamInterval := s.anthropicNativeStreamInterval() + streamInterval := anthropicNativeStreamInterval(s.cfg) pump := newAnthropicNativeLinePump(scanner, streamInterval) defer pump.stop() @@ -328,7 +328,6 @@ func (s *OpenAIGatewayService) handleCCStreamingFromNativeAnthropic( var firstTokenMs *int firstChunk := true clientDisconnected := false - var protocolErr error scanner := bufio.NewScanner(resp.Body) maxLineSize := defaultMaxLineSize @@ -357,7 +356,7 @@ func (s *OpenAIGatewayService) handleCCStreamingFromNativeAnthropic( // 读间隔上限：上游挂住 SSE（不发数据也不断连）时结束排水。上游 ctx 为 // WithoutCancel 且 http.Client 无整体 Timeout，无此界限则客户端断开后 // scanner.Scan() 永久阻塞（见 anthropic native pump 文件注释）。 - streamInterval := s.anthropicNativeStreamInterval() + streamInterval := anthropicNativeStreamInterval(s.cfg) pump := newAnthropicNativeLinePump(scanner, streamInterval) defer pump.stop() @@ -373,9 +372,6 @@ func (s *OpenAIGatewayService) handleCCStreamingFromNativeAnthropic( // 返回——与 messages 主路径 "stream usage incomplete after timeout" 同语义。 onIdle := func() (*OpenAIForwardResult, error) { _ = resp.Body.Close() - if protocolErr != nil { - return resultWithUsage(), protocolErr - } if !clientDisconnected { logger.L().Warn("openai cc via native anthropic stream: data interval timeout", zap.String("request_id", requestID), @@ -401,27 +397,6 @@ func (s *OpenAIGatewayService) handleCCStreamingFromNativeAnthropic( return false } - writeResponseEvent := func(event *apicompat.ResponsesStreamEvent) { - if protocolErr != nil || clientDisconnected { - return - } - chunks := apicompat.ResponsesEventToChatChunks(event, ccState) - if ccState.ProtocolError != "" { - protocolErr = errors.New(ccState.ProtocolError) - MarkResponseCommitted(c) - if !clientDisconnected { - if _, err := fmt.Fprint(c.Writer, buildChatStreamErrorSSE("upstream_protocol_error", ccState.ProtocolError)+"data: [DONE]\n\n"); err != nil { - clientDisconnected = true - } - c.Writer.Flush() - } - return - } - for _, chunk := range chunks { - writeChunk(chunk) - } - } - processAnthropicEvent := func(event *apicompat.AnthropicStreamEvent) bool { if firstChunk { firstChunk = false @@ -442,14 +417,18 @@ func (s *OpenAIGatewayService) handleCCStreamingFromNativeAnthropic( syncAnthropicResponsesUsage(anthState, usage) normalizeAnthropicEventUsageForResponses(event, usage) - // 客户端已断开：跳过转换与写出，继续读上游直到流结束（usage 完整、 // 连接及时归还），不再提前 return。 - if clientDisconnected || protocolErr != nil { + // 客户端已断开：跳过转换与写出，继续读上游直到流结束（usage 完整、 + // 连接及时归还），不再提前 return。 + if clientDisconnected { return false } responsesEvents := apicompat.AnthropicEventToResponsesEvents(event, anthState) for _, resEvt := range responsesEvents { - writeResponseEvent(&resEvt) + ccChunks := apicompat.ResponsesEventToChatChunks(&resEvt, ccState) + for _, chunk := range ccChunks { + writeChunk(chunk) + } } if len(responsesEvents) > 0 { c.Writer.Flush() @@ -494,17 +473,17 @@ func (s *OpenAIGatewayService) handleCCStreamingFromNativeAnthropic( ccState.IncludeUsage = ccState.IncludeUsage || anthropicChatStreamHasUsage(&event, payload) if processAnthropicEvent(&event) { - return resultWithUsage(), protocolErr + return resultWithUsage(), nil } } // Finalize both state machines（客户端已断开时仍执行，保证 usage 汇总完整）。 finalResEvents := apicompat.FinalizeAnthropicResponsesStream(anthState) for _, resEvt := range finalResEvents { - writeResponseEvent(&resEvt) - } - if protocolErr != nil { - return resultWithUsage(), protocolErr + ccChunks := apicompat.ResponsesEventToChatChunks(&resEvt, ccState) + for _, chunk := range ccChunks { + writeChunk(chunk) //nolint:errcheck + } } finalCCChunks := apicompat.FinalizeResponsesChatStream(ccState) for _, chunk := range finalCCChunks { diff --git a/backend/internal/service/openai_gateway_deepseek_input_image_test.go b/backend/internal/service/openai_gateway_deepseek_input_image_test.go index 61f75c193922..b95d75cf3742 100644 --- a/backend/internal/service/openai_gateway_deepseek_input_image_test.go +++ b/backend/internal/service/openai_gateway_deepseek_input_image_test.go @@ -63,7 +63,7 @@ func TestNormalizeDeepSeekResponsesRequestBodyAliasesInputImageURL(t *testing.T) got := normalizeDeepSeekResponsesRequestBody(mapped, body) require.Equal(t, "input_image", gjson.GetBytes(got, "input.0.content.1.type").String()) require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.0.content.1.image_url").String()) - require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.0.content.1.url").String()) + require.False(t, gjson.GetBytes(got, "input.0.content.1.url").Exists()) require.Equal(t, "what is this?", gjson.GetBytes(got, "input.0.content.0.text").String()) require.True(t, gjson.GetBytes(got, "store").Bool(), "mapped openai 账号不应被无状态适配改掉 store") }) @@ -73,7 +73,7 @@ func TestNormalizeDeepSeekResponsesRequestBodyAliasesInputImageURL(t *testing.T) got := normalizeDeepSeekResponsesRequestBody(native, body) require.Equal(t, gjson.String, gjson.GetBytes(got, "input.0.content.1.image_url").Type) require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.0.content.1.image_url").String()) - require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.0.content.1.url").String()) + require.False(t, gjson.GetBytes(got, "input.0.content.1.url").Exists()) }) t.Run("chat_completions_image_url_part_converted", func(t *testing.T) { @@ -81,7 +81,7 @@ func TestNormalizeDeepSeekResponsesRequestBodyAliasesInputImageURL(t *testing.T) got := normalizeDeepSeekResponsesRequestBody(mapped, body) require.Equal(t, "input_image", gjson.GetBytes(got, "input.0.content.1.type").String()) require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.0.content.1.image_url").String()) - require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.0.content.1.url").String()) + require.False(t, gjson.GetBytes(got, "input.0.content.1.url").Exists()) }) t.Run("anthropic_source_becomes_data_uri", func(t *testing.T) { @@ -89,16 +89,24 @@ func TestNormalizeDeepSeekResponsesRequestBodyAliasesInputImageURL(t *testing.T) got := normalizeDeepSeekResponsesRequestBody(native, body) require.Equal(t, "input_image", gjson.GetBytes(got, "input.0.content.1.type").String()) require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.0.content.1.image_url").String()) - require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.0.content.1.url").String()) + require.False(t, gjson.GetBytes(got, "input.0.content.1.url").Exists()) }) - t.Run("lifted_tool_output_image_gets_url", func(t *testing.T) { + t.Run("lifted_tool_output_image_keeps_image_url_only", func(t *testing.T) { body := []byte(`{"model":"deepseek-flash","input":[{"type":"function_call","call_id":"call_image","name":"view_image","arguments":"{}"},{"type":"function_call_output","call_id":"call_image","output":[{"type":"input_image","image_url":"` + deepSeekInputImageDataURI + `"}]}]}`) got := normalizeDeepSeekResponsesRequestBody(mapped, body) require.Equal(t, gjson.String, gjson.GetBytes(got, "input.1.output").Type) require.Equal(t, "input_image", gjson.GetBytes(got, "input.2.content.1.type").String()) require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.2.content.1.image_url").String()) - require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.2.content.1.url").String()) + require.False(t, gjson.GetBytes(got, "input.2.content.1.url").Exists()) + }) + + t.Run("dual_field_input_drops_duplicate_url", func(t *testing.T) { + body := deepSeekUserImageBody(`{"type":"input_image","image_url":"` + deepSeekInputImageDataURI + `","url":"` + deepSeekInputImageDataURI + `"}`) + got := normalizeDeepSeekResponsesRequestBody(native, body) + require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(got, "input.0.content.1.image_url").String()) + require.False(t, gjson.GetBytes(got, "input.0.content.1.url").Exists()) + require.Equal(t, string(got), string(normalizeDeepSeekResponsesRequestBody(native, got))) }) t.Run("file_id_only_does_not_gain_empty_url", func(t *testing.T) { @@ -144,7 +152,7 @@ func TestForwardResponses_OpenAIMappedDeepSeekAliasesInputImageURL(t *testing.T) require.NotNil(t, result) require.Contains(t, upstream.lastReq.URL.String(), "/responses") require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(upstream.lastBody, "input.0.content.1.image_url").String()) - require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(upstream.lastBody, "input.0.content.1.url").String()) + require.False(t, gjson.GetBytes(upstream.lastBody, "input.0.content.1.url").Exists()) } func TestForwardResponses_NativeDeepSeekAliasesInputImageURL(t *testing.T) { @@ -167,5 +175,5 @@ func TestForwardResponses_NativeDeepSeekAliasesInputImageURL(t *testing.T) { require.NotNil(t, result) require.Equal(t, "https://api.deepseek.com/responses", upstream.lastReq.URL.String()) require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(upstream.lastBody, "input.0.content.1.image_url").String()) - require.Equal(t, deepSeekInputImageDataURI, gjson.GetBytes(upstream.lastBody, "input.0.content.1.url").String()) + require.False(t, gjson.GetBytes(upstream.lastBody, "input.0.content.1.url").Exists()) } diff --git a/backend/internal/service/openai_gateway_passthrough.go b/backend/internal/service/openai_gateway_passthrough.go index 399de5f042c2..3307c6f9750d 100644 --- a/backend/internal/service/openai_gateway_passthrough.go +++ b/backend/internal/service/openai_gateway_passthrough.go @@ -845,9 +845,6 @@ func shouldFailoverOpenAIPassthroughResponse(account *Account, statusCode int, r if isOpenAIRequestBudgetRejection(account, statusCode, responseBody) { return true } - if isOpenAIReportedUpstreamFailure(statusCode, responseBody) { - return true - } if account != nil && account.IsOpenAICompatible() && isOpenAIModelNotSupportedError(statusCode, "", responseBody) { return true } @@ -1151,7 +1148,9 @@ func openAIStreamClientOutputStarted(c *gin.Context, localStarted bool) bool { // Lifecycle metadata and transport heartbeats are not model output. func openAIStreamEventIsMetadata(eventType string) bool { switch strings.TrimSpace(eventType) { - case "response.created", "response.in_progress", "keepalive": + // Queuing and moderation metadata do not commit model output or disarm the + // first-output timeout; a following overload must remain safe to retry. + case "response.created", "response.queued", "response.in_progress", "response.metadata", "keepalive": return true default: return false @@ -1213,7 +1212,7 @@ func openAIStreamAddedEventStartsClientOutput(payload []byte, eventType string) default: return true } - case "response.content_part.added": + case "response.content_part.added", "response.content_part.done": part := gjson.GetBytes(payload, "part") if !part.Exists() || !part.IsObject() { return true @@ -1226,7 +1225,7 @@ func openAIStreamAddedEventStartsClientOutput(payload []byte, eventType string) default: return true } - case "response.reasoning_summary_part.added": + case "response.reasoning_summary_part.added", "response.reasoning_summary_part.done": part := gjson.GetBytes(payload, "part") if !part.Exists() || !part.IsObject() || strings.TrimSpace(part.Get("type").String()) != "summary_text" { return true @@ -1242,6 +1241,7 @@ func openAIStreamDataStartsClientOutput(data, eventType string) bool { if trimmed == "" { return false } + field := "" switch strings.TrimSpace(eventType) { case "response.failed": return false @@ -1253,8 +1253,20 @@ func openAIStreamDataStartsClientOutput(data, eventType string) bool { // （content_policy / invalid_request 等）维持原样转发，保留上游错误细节。 payload := []byte(trimmed) return !openAIStreamFailedEventShouldFailover(payload, extractOpenAISSEErrorMessage(payload)) - case "response.output_item.added", "response.content_part.added", "response.reasoning_summary_part.added": + case "response.output_item.added", "response.content_part.added", "response.reasoning_summary_part.added", + "response.content_part.done", "response.reasoning_summary_part.done": return openAIStreamAddedEventStartsClientOutput([]byte(trimmed), eventType) + case "response.output_text.delta", "response.reasoning_summary_text.delta", "response.reasoning_text.delta", + "response.function_call_arguments.delta", "response.custom_tool_call_input.delta", "response.refusal.delta": + field = "delta" + case "response.output_text.done", "response.reasoning_summary_text.done", "response.reasoning_text.done": + field = "text" + case "response.refusal.done": + field = "refusal" + } + if field != "" && gjson.Valid(trimmed) { + value := gjson.Get(trimmed, field) + return value.Type != gjson.String || value.String() != "" } return !openAIStreamEventIsMetadata(eventType) } diff --git a/backend/internal/service/openai_gateway_request_body.go b/backend/internal/service/openai_gateway_request_body.go index c356a92389b4..78feece749e9 100644 --- a/backend/internal/service/openai_gateway_request_body.go +++ b/backend/internal/service/openai_gateway_request_body.go @@ -184,9 +184,8 @@ func deleteOpenAIResponsesNoneReasoningEffortFromObject(account *Account, body m // 强制 store=false 并清除 previous_response_id（DeepSeek / Kimi 官方 // Responses 均不支持服务端状态存储，携带这些字段会被拒绝）。 // -// DeepSeek 另需把 Codex/OpenAI 的 input_image.image_url 写成线上 serde -// 要求的 url 字段；openai 平台但 base_url 指向 api.deepseek.com 的映射 -// 账号同样走这条出站改写（Codex 贴图会 422 missing field url）。 +// DeepSeek 另需把兼容格式的图片输入归一化为官方 image_url 字段；openai +// 平台但 base_url 指向 api.deepseek.com 的映射账号同样走这条出站改写。 // 非 CN / 非 DeepSeek 上游原样返回。 func normalizeDeepSeekResponsesRequestBody(account *Account, body []byte) []byte { if account == nil { @@ -245,10 +244,9 @@ func shouldAliasDeepSeekResponsesInputImages(account *Account) bool { return targetsDeepSeekAPIHost(account) } -// aliasDeepSeekResponsesInputImages 把图片 part 改写成 DeepSeek 能反序列化 -// 的形状：type=input_image，同时带字符串字段 image_url 与 url。 -// 文档与部分 400 提 image_url；线上 serde 要求 url。只发其中一个都会踩坑。 -// 仅有 file_id、没有可用 URL 的 part 原样保留。 +// aliasDeepSeekResponsesInputImages 把图片 part 改写成 DeepSeek 官方形状： +// type=input_image，同时只保留字符串字段 image_url。仅有 file_id、没有可用 +// URL 的 part 原样保留。 func aliasDeepSeekResponsesInputImages(input any) (any, bool) { items, ok := asDeepSeekResponsesSlice(input) if !ok { @@ -330,18 +328,18 @@ func aliasDeepSeekResponsesImagePart(part map[string]any) bool { part["image_url"] = imageURL changed = true } - if current, ok := part["url"].(string); !ok || strings.TrimSpace(current) != imageURL { - part["url"] = imageURL + if _, exists := part["url"]; exists { + delete(part, "url") changed = true } return changed } func extractDeepSeekResponsesImageURL(part map[string]any) string { - if imageURL := deepSeekResponsesURLValue(part["url"]); imageURL != "" { + if imageURL := deepSeekResponsesURLValue(part["image_url"]); imageURL != "" { return imageURL } - if imageURL := deepSeekResponsesURLValue(part["image_url"]); imageURL != "" { + if imageURL := deepSeekResponsesURLValue(part["url"]); imageURL != "" { return imageURL } if imageURL := deepSeekResponsesURLValue(part["image"]); imageURL != "" { diff --git a/backend/internal/service/openai_gateway_responses_anthropic_native.go b/backend/internal/service/openai_gateway_responses_anthropic_native.go index 7181e080db09..bcbde2b7700b 100644 --- a/backend/internal/service/openai_gateway_responses_anthropic_native.go +++ b/backend/internal/service/openai_gateway_responses_anthropic_native.go @@ -171,7 +171,7 @@ func (s *OpenAIGatewayService) handleResponsesBufferedFromNativeAnthropic( var usage ClaudeUsage // 读间隔上限：上游挂住 SSE 时中止组装（缓冲路径尚未提交响应头，可回 502）。 - streamInterval := s.anthropicNativeStreamInterval() + streamInterval := anthropicNativeStreamInterval(s.cfg) pump := newAnthropicNativeLinePump(scanner, streamInterval) defer pump.stop() @@ -367,7 +367,7 @@ func (s *OpenAIGatewayService) handleResponsesStreamingFromNativeAnthropic( // 读间隔上限：上游挂住 SSE（不发数据也不断连）时结束转换循环。上游 ctx 为 // WithoutCancel 且 http.Client 无整体 Timeout，无此界限则 scanner.Scan() // 永久阻塞（见 anthropic native pump 文件注释）。 - streamInterval := s.anthropicNativeStreamInterval() + streamInterval := anthropicNativeStreamInterval(s.cfg) pump := newAnthropicNativeLinePump(scanner, streamInterval) defer pump.stop() diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 32d4b451772c..255579f83291 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -34,7 +34,7 @@ const ( openaiPlatformAPIURL = "https://api.openai.com/v1/responses" openaiPlatformAPIInputTokensURL = "https://api.openai.com/v1/responses/input_tokens" openaiStickySessionTTL = time.Hour // 粘性会话TTL - // 与真实 Codex TUI 的 User-Agent 结构对齐： + // 与真实 Codex CLI 的 User-Agent 结构对齐： // {originator}/{version} ({OS} {OS_version}; {arch}) {terminal} // 缺少 OS/架构/终端后缀的形态易被上游指纹识别为非官方客户端。 // 该后缀是 UA 形态的唯一定义处，buildCodexCLIUserAgent 按运行时版本号复用它。 @@ -58,8 +58,8 @@ const ( openAICompactSessionSeedKey = "openai_compact_session_seed" openAIUpstreamEndpointContextKey = "openai_actual_upstream_endpoint" // codexCLIVersion 是网关对上游声明的 Codex 客户端版本，同时供 codexCLIUserAgent - // 与 version 头使用。上游 /backend-api/codex 在容量紧张时按客户端身份分优先级降载， - // 陈旧版本会被优先丢弃（HTTP 200 + 流内 server_is_overloaded）；非官方客户端配不出 + // 与 version 头使用。上游 /backend-api/codex 会按客户端身份进行风控分层， + // 陈旧版本可能被优先丢弃（HTTP 200 + 流内 server_is_overloaded）；非官方客户端配不出 // 官方身份时整体回退到本常量，因此它必须跟随官方 CLI 的当前发布版本， // 落后多个版本会让这些请求稳定落在被优先丢弃的一侧。 codexCLIVersion = "0.163.0" diff --git a/backend/internal/service/openai_gateway_upstream_errors.go b/backend/internal/service/openai_gateway_upstream_errors.go index c118a65bf91f..ccadfe3ec111 100644 --- a/backend/internal/service/openai_gateway_upstream_errors.go +++ b/backend/internal/service/openai_gateway_upstream_errors.go @@ -7,7 +7,9 @@ import ( "fmt" "io" "net/http" + "strconv" "strings" + "time" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" @@ -190,6 +192,63 @@ func isOpenAIRequestScopedCapacityShed(upstreamMsg string, upstreamBody []byte) (!gjson.ValidBytes(upstreamBody) && isOpenAICapacityShedMessage(string(upstreamBody))) } +const openAIRequestScopedRetryMaxDelay = 8 * time.Second + +func openAIRequestScopedRetryDelay(responseHeaders http.Header, responseBody []byte) time.Duration { + type retryAfterValue struct { + raw string + milliseconds bool + } + for _, value := range []retryAfterValue{ + {raw: responseHeaders.Get("Retry-After")}, + {raw: responseHeaders.Get("Retry-After-Ms"), milliseconds: true}, + {raw: gjson.GetBytes(responseBody, "error.headers.retry-after").String()}, + {raw: gjson.GetBytes(responseBody, "response.error.headers.retry-after").String()}, + {raw: gjson.GetBytes(responseBody, "error.headers.retry-after-ms").String(), milliseconds: true}, + {raw: gjson.GetBytes(responseBody, "response.error.headers.retry-after-ms").String(), milliseconds: true}, + } { + raw := strings.TrimSpace(value.raw) + if raw == "" { + continue + } + var delay time.Duration + var ok bool + if value.milliseconds { + delay, ok = parseOpenAIRequestScopedRetryDelayMilliseconds(raw) + } else { + delay, ok = parseOpenAIRequestScopedRetryDelay(raw) + } + if !ok || delay <= 0 { + continue + } + if delay > openAIRequestScopedRetryMaxDelay { + return openAIRequestScopedRetryMaxDelay + } + return delay + } + return 0 +} + +func parseOpenAIRequestScopedRetryDelay(raw string) (time.Duration, bool) { + if seconds, err := strconv.ParseFloat(raw, 64); err == nil && seconds > 0 { + return time.Duration(seconds * float64(time.Second)), true + } + if resetAt, err := http.ParseTime(raw); err == nil { + delay := time.Until(resetAt) + if delay > 0 { + return delay, true + } + } + return 0, false +} + +func parseOpenAIRequestScopedRetryDelayMilliseconds(raw string) (time.Duration, bool) { + if millis, err := strconv.ParseFloat(raw, 64); err == nil && millis > 0 { + return time.Duration(millis * float64(time.Millisecond)), true + } + return 0, false +} + func isOpenAIContextWindowError(upstreamMsg string, upstreamBody []byte) bool { match := func(text string) bool { lower := strings.ToLower(strings.TrimSpace(text)) @@ -275,9 +334,6 @@ func (s *OpenAIGatewayService) shouldFailoverOpenAIUpstreamResponse(account *Acc if isOpenAIModelNotSupportedError(statusCode, upstreamMsg, upstreamBody) { return account != nil && account.IsOpenAICompatible() } - if isOpenAIReportedUpstreamFailure(statusCode, upstreamBody) { - return true - } if isOpenAIHTTPUpstreamAccessStateError(statusCode, upstreamMsg, upstreamBody) { return true } @@ -303,28 +359,6 @@ func (s *OpenAIGatewayService) shouldFailoverOpenAIUpstreamResponse(account *Acc return isOpenAITransientProcessingError(statusCode, upstreamMsg, upstreamBody) } -// A proxy can report its own upstream failure using HTTP 400. The structured -// error type, not the number or message text alone, distinguishes that from a -// client's invalid request. An identified parameter always stays client-scoped. -func isOpenAIReportedUpstreamFailure(status int, body []byte) bool { - if status != http.StatusBadRequest || gjson.GetBytes(body, "error.type").String() != "upstream_error" || - strings.TrimSpace(gjson.GetBytes(body, "error.param").String()) != "" { - return false - } - if isOpenAIContinuationStateError("", body) || isOpenAIContextWindowError("", body) || isOpenAIInstructionsRequiredError(status, "", body) { - return false - } - if hit, _, _ := detectOpenAICyberPolicy(body); hit { - return false - } - switch gjson.GetBytes(body, "error.code").String() { - case "", "400", "upstream_error", "server_error": - return true - default: - return false - } -} - func isOpenAIRequestBudgetRejection(account *Account, status int, body []byte) bool { return account != nil && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey && (status == http.StatusForbidden || status == http.StatusPaymentRequired) && @@ -486,8 +520,7 @@ func classifyOpenAIRequestRejection(statusCode int, upstreamMsg string, upstream isOpenAIRequestScopedSafetyRejection(upstreamBody) || isOpenAICompatibleModelNotFound400(upstreamBody) || isOpenAIModelNotSupportedError(statusCode, upstreamMsg, upstreamBody) || - isOpenAIHTTPUpstreamAccessStateError(statusCode, upstreamMsg, upstreamBody) || - isOpenAIReportedUpstreamFailure(statusCode, upstreamBody) { + isOpenAIHTTPUpstreamAccessStateError(statusCode, upstreamMsg, upstreamBody) { return "" } @@ -574,14 +607,8 @@ func newOpenAIUpstreamFailoverError( RetryableOnSameAccount: retryableOnSameAccount || requestScopedCapacity, RequestScopedTransient: requestScopedCapacity, } - if isOpenAIReportedUpstreamFailure(statusCode, responseBody) { - failoverErr.RetryableOnSameAccount = false - failoverErr.Scope = GatewayFailureScopeRequest - failoverErr.SuppressAccountHealthPenalty = true - failoverErr.Reason = GatewayFailureReason("openai_upstream_reported_failure") - failoverErr.ClientStatusCode = http.StatusBadGateway - failoverErr.ClientErrorCode = "upstream_request_failed" - failoverErr.ClientMessage = "Upstream failed to process the request" + if requestScopedCapacity { + failoverErr.SameAccountRetryDelay = openAIRequestScopedRetryDelay(responseHeaders, responseBody) } if isOpenAIRequestBodyTooLargeError(statusCode, upstreamMsg, responseBody) { failoverErr.RetryableOnSameAccount = false diff --git a/backend/internal/service/openai_gateway_upstream_errors_test.go b/backend/internal/service/openai_gateway_upstream_errors_test.go index 30d8fa8fa891..41c622139dc2 100644 --- a/backend/internal/service/openai_gateway_upstream_errors_test.go +++ b/backend/internal/service/openai_gateway_upstream_errors_test.go @@ -33,7 +33,7 @@ func TestClassifyOpenAIRequestRejection(t *testing.T) { {"model missing without code", 400, `{"error":{"message":"unknown provider for model missing-model"}}`, ""}, {"model unavailable", 400, `{"error":{"type":"model_not_supported","code":400,"message":"model is not supported"}}`, ""}, {"workspace access", 400, `{"error":{"code":"deactivated_workspace"}}`, ""}, - {"reported upstream failure", 400, `{"error":{"type":"upstream_error","message":"failed"}}`, ""}, + {"upstream error envelope", 400, `{"error":{"type":"upstream_error","message":"failed"}}`, "unclassified_bad_request"}, {"authentication HTTP error", 401, `{"error":{"type":"invalid_request_error"}}`, ""}, {"upstream gateway error", 502, `{"error":{"message":"generic compatibility failure"}}`, ""}, } diff --git a/backend/internal/service/openai_oauth_passthrough_test.go b/backend/internal/service/openai_oauth_passthrough_test.go index 0724cee242d4..721e05bc5030 100644 --- a/backend/internal/service/openai_oauth_passthrough_test.go +++ b/backend/internal/service/openai_oauth_passthrough_test.go @@ -1299,7 +1299,7 @@ func TestOpenAIGatewayService_OAuthLegacy_CompositeCodexUAUsesCodexOriginator(t _, err := svc.Forward(context.Background(), c, account, inputBody) require.NoError(t, err) require.NotNil(t, upstream.lastReq) - // 浏览器型复合 UA 被替换为该账号的 Codex TUI 身份 UA， + // 浏览器型复合 UA 被替换为该账号固定设备的 Codex CLI 身份 UA， // originator 随最终 UA 配套（issue #3901）。 require.Equal(t, resolveCodexOutboundIdentityForAccount(account, "").userAgent, upstream.lastReq.Header.Get("User-Agent")) require.Equal(t, openai.CodexDefaultOriginator, upstream.lastReq.Header.Get("originator")) diff --git a/backend/internal/service/openai_setup_token_compat_test.go b/backend/internal/service/openai_setup_token_compat_test.go index 7aa52223b3e0..85ba6394ba78 100644 --- a/backend/internal/service/openai_setup_token_compat_test.go +++ b/backend/internal/service/openai_setup_token_compat_test.go @@ -11,6 +11,7 @@ import ( "testing" "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/openai" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" @@ -212,7 +213,7 @@ func TestOpenAISetupTokenMessagesUsesCodexBridgeAndTurnState(t *testing.T) { require.Equal(t, chatgptCodexURL, upstream.requests[0].URL.String()) require.Equal(t, "Bearer setup-token-value", upstream.requests[0].Header.Get("Authorization")) require.Equal(t, "chatgpt-setup", upstream.requests[0].Header.Get("chatgpt-account-id")) - requireOpenAIMessagesCodexIdentity(t, upstream.requests[0], resolveCodexOutboundIdentityForAccount(account, "").userAgent, "codex-tui") + requireOpenAIMessagesCodexIdentity(t, upstream.requests[0], resolveCodexOutboundIdentityForAccount(account, "").userAgent, openai.CodexCLIOriginator) require.Empty(t, upstream.requests[0].Header.Get("x-codex-turn-state")) secondBody := []byte(`{"model":"claude-sonnet-4-5","max_tokens":16,"messages":[{"role":"user","content":"next"}],"stream":false}`) @@ -229,7 +230,7 @@ func TestOpenAISetupTokenMessagesUsesCodexBridgeAndTurnState(t *testing.T) { require.Equal(t, "turn_state_setup", upstream.requests[1].Header.Get("x-codex-turn-state")) require.Equal(t, generateSessionUUID(isolateOpenAIUpstreamSessionID(0, account, "stable-cache-key")), upstream.requests[1].Header.Get("session-id")) require.Empty(t, upstream.requests[1].Header.Get("conversation_id")) - requireOpenAIMessagesCodexIdentity(t, upstream.requests[1], resolveCodexOutboundIdentityForAccount(account, "").userAgent, "codex-tui") + requireOpenAIMessagesCodexIdentity(t, upstream.requests[1], resolveCodexOutboundIdentityForAccount(account, "").userAgent, openai.CodexCLIOriginator) } func openAISetupTokenCompatAccount(id int64) *Account { diff --git a/backend/internal/service/openai_ws_forwarder_success_test.go b/backend/internal/service/openai_ws_forwarder_success_test.go index 866475f4301b..389bb1e2e808 100644 --- a/backend/internal/service/openai_ws_forwarder_success_test.go +++ b/backend/internal/service/openai_ws_forwarder_success_test.go @@ -1024,9 +1024,9 @@ func TestOpenAIGatewayService_Forward_WSv2_OAuthHonorsAccountUserAgent(t *testin result, err := svc.Forward(context.Background(), c, account, body) require.NoError(t, err) require.NotNil(t, result) - require.Equal(t, "codex-tui", captureDialer.lastHeaders.Get("originator")) + require.Equal(t, openai.CodexCLIOriginator, captureDialer.lastHeaders.Get("originator")) require.Equal(t, - "codex-tui/"+codexCLIVersion+" (Mac OS X 15.1.0; arm64) iTerm.app", + "codex_cli_rs/"+codexCLIVersion+" (Mac OS X 15.1.0; arm64) iTerm.app", captureDialer.lastHeaders.Get("user-agent"), ) require.Equal(t, codexCLIVersion, captureDialer.lastHeaders.Get("version")) diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index 39183a868d17..bcd6e8d91a4a 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -1300,6 +1300,10 @@ func TestProxyOpenAIWSHTTPBridgeTurnStagesMetadataAndRelaysKeepaliveBeforeCapaci preamble := strings.Join([]string{ `data: {"type":"response.created","response":{"id":"resp_shed"}}`, "", + `data: {"type":"response.queued","response":{"id":"resp_shed","status":"queued"}}`, + "", + `data: {"type":"response.metadata","response_id":"resp_shed","metadata":{"openai_chatgpt_moderation_metadata":{"prompt":{}}}}`, + "", `data: {"type":"response.in_progress","response":{"id":"resp_shed"}}`, "", `data: {"type":"keepalive"}`, diff --git a/backend/internal/service/ops_upstream_context.go b/backend/internal/service/ops_upstream_context.go index cf67c6cc0d00..1defdb4f77dd 100644 --- a/backend/internal/service/ops_upstream_context.go +++ b/backend/internal/service/ops_upstream_context.go @@ -411,8 +411,7 @@ func IsGatewayFixedClientMessage(message string) bool { openAISilentRefusalClientMessage, openAIUpstreamAccessUnavailableClientMessage, geminiCustomCodeSkippedClientMessage, - "Temporary upstream failure", - "Upstream failed to process the request": + "Temporary upstream failure": return true } return false diff --git a/backend/internal/service/pricing_service_test.go b/backend/internal/service/pricing_service_test.go index 08f5171a7c08..d72e2f121414 100644 --- a/backend/internal/service/pricing_service_test.go +++ b/backend/internal/service/pricing_service_test.go @@ -819,6 +819,43 @@ func TestDefaultPricingIncludesGemini4ArgonRates(t *testing.T) { require.Contains(t, pricingSvc.ListModelNamesByProvider("gemini"), "gemini-4-argon") } +func TestDefaultPricingIncludesGrok47FastRates(t *testing.T) { + data, err := os.ReadFile(filepath.Join("..", "..", "resources", "model-pricing", "model_prices_and_context_window.json")) + require.NoError(t, err) + + pricingSvc := &PricingService{} + pricingData, err := pricingSvc.parsePricingData(data) + require.NoError(t, err) + pricingSvc.pricingData = pricingData + + for _, model := range []string{"grok-4.7-fast", "xai/grok-4.7-fast"} { + t.Run(model, func(t *testing.T) { + pricing := pricingSvc.GetModelPricing(model) + require.NotNil(t, pricing) + require.Equal(t, "xai", pricing.LiteLLMProvider) + require.InDelta(t, 4e-6, pricing.InputCostPerToken, 1e-15) + require.InDelta(t, 12e-6, pricing.OutputCostPerToken, 1e-15) + require.InDelta(t, 1e-6, pricing.CacheReadInputTokenCost, 1e-15) + // Long-context Fast is $6 / $1.50 / $18, i.e. 1.5x the Fast base card. + require.Equal(t, 200000, pricing.LongContextInputTokenThreshold) + require.InDelta(t, 1.5, pricing.LongContextInputCostMultiplier, 1e-12) + require.InDelta(t, 1.5, pricing.LongContextOutputCostMultiplier, 1e-12) + }) + } + // The Fast entry must not shadow the standard grok-4.7 card. + standard := pricingSvc.GetModelPricing("grok-4.7") + require.NotNil(t, standard) + require.InDelta(t, 2e-6, standard.InputCostPerToken, 1e-15) + + billing := NewBillingService(&config.Config{}, pricingSvc) + cost, err := billing.CalculateCost("grok-4.7-fast", UsageTokens{InputTokens: 150_000, CacheReadTokens: 50_000, OutputTokens: 10_000}, 1) + require.NoError(t, err) + require.True(t, cost.LongContextBillingApplied) + require.InDelta(t, 0.9, cost.InputCost, 1e-12) + require.InDelta(t, 0.075, cost.CacheReadCost, 1e-12) + require.InDelta(t, 0.18, cost.OutputCost, 1e-12) +} + func TestDefaultPricingUsesCurrentCodexAutoReviewBaseRates(t *testing.T) { data, err := os.ReadFile(filepath.Join("..", "..", "resources", "model-pricing", "model_prices_and_context_window.json")) require.NoError(t, err) diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index d22c453434fa..c6b6ea2ca4ac 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -1046,10 +1046,19 @@ func buildForbiddenErrorMessage(prefix string, upstreamMsg string, responseBody return prefix + fallback } -// handle403 处理 403 Forbidden 错误 -// Antigravity 平台区分 validation/violation/generic 三种类型，均 SetError 永久禁用； -// 其他平台保持原有 SetError 行为。 +// handle403 处理 403 Forbidden 错误。Cloudflare 1010 是请求/边缘级拦截， +// 不应将其视为账号级故障；其他真实 403 继续按平台策略处理。 func (s *RateLimitService) handle403(ctx context.Context, account *Account, upstreamMsg string, responseBody []byte) (shouldDisable bool) { + if isCloudflareBotBlockResponse(responseBody) { + slog.Warn( + "403_cloudflare_bot_block_skips_account_penalty", + "account_id", account.ID, + "platform", account.Platform, + "upstream_message", upstreamMsg, + ) + return false + } + if account.Platform == PlatformAntigravity { return s.handleAntigravity403(ctx, account, upstreamMsg, responseBody) } @@ -1107,16 +1116,6 @@ func (s *RateLimitService) handleOpenAI403(ctx context.Context, account *Account ) return false } - if isCloudflareBotBlockResponse(responseBody) { - slog.Warn( - "openai_403_cloudflare_bot_block_skips_account_penalty", - "account_id", account.ID, - "platform", account.Platform, - "upstream_message", upstreamMsg, - ) - return false - } - msg := buildForbiddenErrorMessage( "Access forbidden (403):", upstreamMsg, @@ -1166,9 +1165,62 @@ func (s *RateLimitService) handleOpenAI403(ctx context.Context, account *Account // request/edge-level failure and must not consume the account 403 strike budget. func isCloudflareBotBlockResponse(body []byte) bool { normalized := strings.ToLower(strings.TrimSpace(string(body))) + var payload any + if json.Unmarshal(body, &payload) == nil { + return containsCloudflareBotBlockCode(payload) + } return strings.Contains(normalized, "error code: 1010") } +func containsCloudflareBotBlockCode(value any) bool { + switch node := value.(type) { + case map[string]any: + cloudflareError := false + for key, child := range node { + if strings.EqualFold(key, "cloudflare_error") { + if is, ok := child.(bool); ok && is { + cloudflareError = true + } + } + } + for key, child := range node { + if strings.EqualFold(key, "error_name") { + if name, ok := child.(string); ok && strings.EqualFold(strings.TrimSpace(name), "browser_signature_banned") { + return true + } + } + if cloudflareError && (strings.EqualFold(key, "code") || strings.EqualFold(key, "error_code")) { + if isCloudflareBotBlockCode(child) { + return true + } + } + if containsCloudflareBotBlockCode(child) { + return true + } + } + case []any: + for _, child := range node { + if containsCloudflareBotBlockCode(child) { + return true + } + } + } + return false +} + +func isCloudflareBotBlockCode(value any) bool { + switch code := value.(type) { + case float64: + return code == 1010 + case json.Number: + return code == "1010" + case string: + return strings.TrimSpace(code) == "1010" + default: + return false + } +} + // handleAntigravity403 处理 Antigravity 平台的 403 错误 // validation（需要验证）→ 永久 SetError（需人工去 Google 验证后恢复） // violation（违规封号）→ 永久 SetError（需人工处理） diff --git a/backend/internal/service/ratelimit_service_403_html_test.go b/backend/internal/service/ratelimit_service_403_html_test.go index 8768b41fea61..8048dd3cabe8 100644 --- a/backend/internal/service/ratelimit_service_403_html_test.go +++ b/backend/internal/service/ratelimit_service_403_html_test.go @@ -104,8 +104,8 @@ func TestHandleUpstreamError_OpenAIHTML403RepeatedNeverEscalates(t *testing.T) { h.requireNoAccountPenalty(t) } -func TestHandleUpstreamError_OpenAICloudflare1010DoesNotPenalizeAccount(t *testing.T) { - for _, platform := range []string{PlatformOpenAI, PlatformOpenCodeGo} { +func TestHandleUpstreamError_Cloudflare1010DoesNotPenalizeAccount(t *testing.T) { + for _, platform := range []string{PlatformOpenAI, PlatformOpenCodeGo, PlatformAnthropic, PlatformGemini, PlatformAntigravity} { t.Run(platform, func(t *testing.T) { repo := &rateLimitAccountRepoStub{} counter := &countingOpenAI403CounterCache{ @@ -130,6 +130,67 @@ func TestHandleUpstreamError_OpenAICloudflare1010DoesNotPenalizeAccount(t *testi } } +func TestHandleUpstreamError_Cloudflare1010JSONDoesNotPenalizeAccount(t *testing.T) { + for _, body := range []string{ + `{"cloudflare_error":true,"error_code":1010,"error_name":"browser_signature_banned","detail":"blocked"}`, + `{"error_name":"browser_signature_banned","detail":"blocked by Cloudflare"}`, + } { + t.Run(body, func(t *testing.T) { + repo := &rateLimitAccountRepoStub{} + counter := &countingOpenAI403CounterCache{ + openAI403CounterCacheStub: openAI403CounterCacheStub{counts: []int64{openAI403DisableThreshold}}, + } + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + svc.SetOpenAI403CounterCache(counter) + account := &Account{ID: 602, Platform: PlatformAnthropic, Type: AccountTypeAPIKey} + + require.False(t, svc.HandleUpstreamError( + context.Background(), account, http.StatusForbidden, http.Header{}, []byte(body), + )) + require.Equal(t, 0, counter.increments) + require.Equal(t, 0, repo.setErrorCalls) + require.Equal(t, 0, repo.tempCalls) + }) + } +} + +func TestIsCloudflareBotBlockResponseOnlyMatches1010(t *testing.T) { + cases := []struct { + name string + body string + want bool + }{ + {"plain text", "error code: 1010", true}, + {"root numeric JSON code", `{"cloudflare_error":true,"code":1010}`, true}, + {"nested string JSON code", `{"errors":[{"cloudflare_error":true,"code":"1010"}]}`, true}, + {"cloudflare error code", `{"cloudflare_error":true,"error_code":1010}`, true}, + {"browser signature name", `{"error_name":"browser_signature_banned"}`, true}, + {"different JSON code", `{"code":403,"message":"forbidden"}`, false}, + {"unrelated code 1010", `{"error":{"code":1010,"message":"account banned"}}`, false}, + {"message mentions 1010", `{"message":"error code: 1010"}`, false}, + {"malformed JSON", `{"code":1010`, false}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.want, isCloudflareBotBlockResponse([]byte(tc.body))) + }) + } +} + +func TestHandleUpstreamError_AnthropicGenuine403StillPenalizes(t *testing.T) { + repo := &rateLimitAccountRepoStub{} + svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + account := &Account{ID: 603, Platform: PlatformAnthropic, Type: AccountTypeAPIKey} + + require.True(t, svc.HandleUpstreamError( + context.Background(), account, http.StatusForbidden, http.Header{}, + []byte(`{"error":{"type":"permission_error","message":"invalid API key"}}`), + )) + require.Equal(t, 1, repo.setErrorCalls) + require.Contains(t, repo.lastErrorMsg, "invalid API key") +} + // 对照不变式：真正的结构化 JSON 403 是账号级证据，处罚链路必须原样保留。 // 缺了这组断言，上面的跳过逻辑一旦写宽就会把真实的封号 403 也放过去。 func TestHandleUpstreamError_OpenAIStructured403StillPenalizes(t *testing.T) { diff --git a/backend/internal/service/request_metadata.go b/backend/internal/service/request_metadata.go index 5c81bbf1220c..4a521c20f610 100644 --- a/backend/internal/service/request_metadata.go +++ b/backend/internal/service/request_metadata.go @@ -2,6 +2,7 @@ package service import ( "context" + "sync" "sync/atomic" "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" @@ -11,6 +12,15 @@ type requestMetadataContextKey struct{} var requestMetadataKey = requestMetadataContextKey{} +type singleAccountRetryProbeContextKey struct{} + +// singleAccountRetryProbe 在首次需要时判定是否进入单账号退避重试模式，同一请求内至多判定一次。 +type singleAccountRetryProbe struct { + once sync.Once + check func(ctx context.Context) bool + result bool +} + type RequestMetadata struct { IsMaxTokensOneHaikuRequest *bool ThinkingEnabled *bool @@ -107,6 +117,29 @@ func WithSingleAccountRetry(ctx context.Context, value bool, bridgeOldKeys bool) }) } +// WithSingleAccountRetryProbe 挂载单账号退避重试的延迟判定；显式设置的 SingleAccountRetry 标记优先于它。 +func WithSingleAccountRetryProbe(ctx context.Context, check func(ctx context.Context) bool) context.Context { + if ctx == nil || check == nil { + return ctx + } + return context.WithValue(ctx, singleAccountRetryProbeContextKey{}, &singleAccountRetryProbe{check: check}) +} + +func singleAccountRetryProbeResult(ctx context.Context) bool { + if ctx == nil { + return false + } + probe, _ := ctx.Value(singleAccountRetryProbeContextKey{}).(*singleAccountRetryProbe) + if probe == nil { + return false + } + // 判定结果在整个请求内复用，不能因调用时 ctx 已取消而被记成 false。 + probe.once.Do(func() { + probe.result = probe.check(context.WithoutCancel(ctx)) + }) + return probe.result +} + func WithAccountSwitchCount(ctx context.Context, value int, bridgeOldKeys bool) context.Context { return updateRequestMetadata(ctx, bridgeOldKeys, func(md *RequestMetadata) { v := value diff --git a/backend/internal/service/responses_chat_protocol_callers_test.go b/backend/internal/service/responses_chat_protocol_callers_test.go index 8cfc4607e5be..5885217216bf 100644 --- a/backend/internal/service/responses_chat_protocol_callers_test.go +++ b/backend/internal/service/responses_chat_protocol_callers_test.go @@ -1,262 +1,11 @@ package service import ( - "encoding/json" - "errors" - "io" - "net/http" - "net/http/httptest" - "strings" "testing" - "time" - "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" - "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) -type protocolChatWire struct { - finishes []string - tools map[int]apicompat.ChatToolCall - usages []apicompat.ChatUsage - errors int - done int -} - -func readProtocolChatWire(t *testing.T, body string) protocolChatWire { - t.Helper() - out := protocolChatWire{tools: make(map[int]apicompat.ChatToolCall)} - for _, line := range strings.Split(body, "\n") { - payload, ok := strings.CutPrefix(line, "data: ") - if !ok { - continue - } - if payload == "[DONE]" { - out.done++ - continue - } - var envelope struct { - Error *struct{ Code, Message string } `json:"error"` - } - require.NoError(t, json.Unmarshal([]byte(payload), &envelope)) - if envelope.Error != nil { - out.errors++ - require.Equal(t, "upstream_protocol_error", envelope.Error.Code) - require.Equal(t, "Upstream Responses function-call stream is incomplete or inconsistent", envelope.Error.Message) - continue - } - var chunk apicompat.ChatCompletionsChunk - require.NoError(t, json.Unmarshal([]byte(payload), &chunk)) - if chunk.Usage != nil { - out.usages = append(out.usages, *chunk.Usage) - } - for _, choice := range chunk.Choices { - if choice.FinishReason != nil { - out.finishes = append(out.finishes, *choice.FinishReason) - } - for _, delta := range choice.Delta.ToolCalls { - require.NotNil(t, delta.Index) - tool := out.tools[*delta.Index] - tool.ID += delta.ID - tool.Function.Name += delta.Function.Name - tool.Function.Arguments += delta.Function.Arguments - out.tools[*delta.Index] = tool - } - } - } - return out -} - -func protocolAnthropicStream(arguments string, closeTool bool) string { - events := []string{ - `{"type":"message_start","message":{"id":"msg_fixture","type":"message","role":"assistant","content":[],"model":"fixture","usage":{"input_tokens":13}}}`, - `{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"call_fixture","name":"lookup","input":{}}}`, - } - delta, _ := json.Marshal(map[string]any{"type": "content_block_delta", "index": 0, "delta": map[string]any{"type": "input_json_delta", "partial_json": arguments}}) - events = append(events, string(delta)) - if closeTool { - events = append(events, - `{"type":"content_block_stop","index":0}`, - `{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":19}}`, - `{"type":"message_stop"}`, - ) - } else { - // The malformed tool is not closed until the real EOF finalizer runs. - events = append(events, `{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":19}}`) - } - var stream strings.Builder - for _, event := range events { - var kind struct{ Type string } - _ = json.Unmarshal([]byte(event), &kind) - _, _ = stream.WriteString("event: " + kind.Type + "\ndata: " + event + "\n\n") - } - return stream.String() -} - -func TestAnthropicChatCallersSurfaceProtocolFailure(t *testing.T) { - gin.SetMode(gin.TestMode) - for _, native := range []bool{false, true} { - name := map[bool]string{false: "anthropic_gateway", true: "native_anthropic"}[native] - for _, mode := range []string{"item_done", "eof_finalize", "valid_control"} { - t.Run(name+"/"+mode, func(t *testing.T) { - args := `{"x":` - if mode == "valid_control" { - args = `{"x":1}` - } - resp := &http.Response{Header: make(http.Header), Body: io.NopCloser(strings.NewReader(protocolAnthropicStream(args, mode != "eof_finalize")))} - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - var input, output int - var err error - if native { - var result *OpenAIForwardResult - result, err = (&OpenAIGatewayService{}).handleCCStreamingFromNativeAnthropic(resp, c, "fixture", "fixture", "fixture", nil, time.Now()) - require.NotNil(t, result) - input, output = result.Usage.InputTokens, result.Usage.OutputTokens - } else { - var result *ForwardResult - result, err = (&GatewayService{}).handleCCStreamingFromAnthropic(resp, c, "fixture", "fixture", nil, time.Now()) - require.NotNil(t, result) - input, output = result.Usage.InputTokens, result.Usage.OutputTokens - } - require.Equal(t, 13, input) - require.Equal(t, 19, output, "usage after the bad item must still be drained") - wire := readProtocolChatWire(t, rec.Body.String()) - require.Equal(t, 1, wire.done) - require.Equal(t, "call_fixture", wire.tools[0].ID) - require.Equal(t, "lookup", wire.tools[0].Function.Name) - require.Equal(t, args, wire.tools[0].Function.Arguments) - if mode == "valid_control" { - require.NoError(t, err) - require.Zero(t, wire.errors) - require.Equal(t, []string{"tool_calls"}, wire.finishes) - require.Len(t, wire.usages, 1) - require.Equal(t, 13, wire.usages[0].PromptTokens) - require.Equal(t, 19, wire.usages[0].CompletionTokens) - } else { - assertProtocolChatFailed(t, c, wire, err) - } - }) - } - } -} - -func assertProtocolChatFailed(t *testing.T, c *gin.Context, wire protocolChatWire, err error) { - t.Helper() - require.Error(t, err, "the actual caller must not report successful delivery") - var failover *UpstreamFailoverError - require.False(t, errors.As(err, &failover), "corruption after output commitment must not request a retry") - require.True(t, IsResponseCommitted(c), "the handler must not append another error response") - require.Equal(t, 1, wire.errors) - require.Equal(t, 1, wire.done) - require.Empty(t, wire.finishes) - require.Empty(t, wire.usages) -} - -func TestNativeAnthropicChatProtocolGuardPreservesDisconnectDrain(t *testing.T) { - gin.SetMode(gin.TestMode) - // The wire is valid, but the client disconnects after the first argument - // prefix. The caller deliberately stops converting and only drains usage. - body := protocolAnthropicStream(`{"x":`, true) - body = strings.Replace(body, "event: content_block_stop\n", "event: content_block_delta\ndata: "+ - `{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"1}"}}`+"\n\nevent: content_block_stop\n", 1) - resp := &http.Response{Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))} - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 2} - result, err := (&OpenAIGatewayService{}).handleCCStreamingFromNativeAnthropic(resp, c, "fixture", "fixture", "fixture", nil, time.Now()) - require.NoError(t, err, "an intentionally unconverted tail is not evidence of upstream corruption") - require.NotNil(t, result) - require.True(t, result.ClientDisconnect) - require.Equal(t, 13, result.Usage.InputTokens) - require.Equal(t, 19, result.Usage.OutputTokens) -} - -func TestGeminiChatCallerSurfacesProtocolFailure(t *testing.T) { - gin.SetMode(gin.TestMode) - for _, mode := range []string{"closed_by_text", "eof_finalize", "valid_control"} { - t.Run(mode, func(t *testing.T) { - args := `{"x":` - if mode == "valid_control" { - args = `{"x":1}` - } - encoded, err := json.Marshal(args) - require.NoError(t, err) - body := `data: {"candidates":[{"content":{"parts":[{"functionCall":{"name":"lookup","args":` + string(encoded) + `}}]}}],"usageMetadata":{"promptTokenCount":13,"candidatesTokenCount":2}}` + "\n\n" - if mode == "closed_by_text" { - body += `data: {"candidates":[{"content":{"parts":[{"text":"after invalid tool"}]}}]}` + "\n\n" - } - body += `data: {"candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":13,"candidatesTokenCount":19}}` + "\n\n" - resp := &http.Response{Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))} - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - result, err := (&GeminiMessagesCompatService{}).handleChatCompletionsStreamingResponseFromGemini(c, resp, time.Now(), "fixture", false, true) - require.NotNil(t, result) - require.Equal(t, 13, result.usage.InputTokens) - require.Equal(t, 19, result.usage.OutputTokens) - wire := readProtocolChatWire(t, rec.Body.String()) - require.Equal(t, args, wire.tools[0].Function.Arguments) - require.Equal(t, "lookup", wire.tools[0].Function.Name) - if mode == "valid_control" { - require.NoError(t, err) - require.Zero(t, wire.errors) - require.Equal(t, 1, wire.done) - require.Equal(t, []string{"tool_calls"}, wire.finishes) - require.Len(t, wire.usages, 1) - require.Equal(t, 13, wire.usages[0].PromptTokens) - require.Equal(t, 19, wire.usages[0].CompletionTokens) - } else { - assertProtocolChatFailed(t, c, wire, err) - } - }) - } -} - -func TestAntigravityChatCallerSurfacesProtocolFailure(t *testing.T) { - gin.SetMode(gin.TestMode) - for _, mode := range []string{"duplicate_call", "read_error_after_failure", "valid_control"} { - t.Run(mode, func(t *testing.T) { - secondID := "call_fixture" - if mode == "valid_control" { - secondID = "call_second" - } - body := `data: {"response":{"responseId":"resp_fixture","candidates":[{"content":{"parts":[{"functionCall":{"id":"call_fixture","name":"lookup","args":{"x":1}}},{"functionCall":{"id":"` + secondID + `","name":"lookup","args":{"x":2}}}]}}],"usageMetadata":{"promptTokenCount":13,"candidatesTokenCount":2}}}` + "\n\n" + - `data: {"response":{"candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":13,"candidatesTokenCount":19}}}` + "\n\n" - reader := io.NopCloser(strings.NewReader(body)) - if mode == "read_error_after_failure" { - reader = &antigravityCompatErrorReader{data: []byte(body), err: io.ErrUnexpectedEOF} - } - resp := &http.Response{Header: make(http.Header), Body: reader} - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - result, err := (&AntigravityGatewayService{}).handleChatCompletionsStreamingFromAntigravity(c, resp, time.Now(), "fixture", true) - require.NotNil(t, result) - require.Equal(t, 13, result.usage.InputTokens) - require.Equal(t, 19, result.usage.OutputTokens) - wire := readProtocolChatWire(t, rec.Body.String()) - require.Equal(t, "call_fixture", wire.tools[0].ID) - if mode == "valid_control" { - require.NoError(t, err) - require.Zero(t, wire.errors) - require.Equal(t, 1, wire.done) - require.Equal(t, []string{"tool_calls"}, wire.finishes) - require.Len(t, wire.usages, 1) - require.Equal(t, 13, wire.usages[0].PromptTokens) - require.Equal(t, 19, wire.usages[0].CompletionTokens) - require.Equal(t, "call_second", wire.tools[1].ID) - require.Equal(t, `{"x":1}`, wire.tools[0].Function.Arguments) - require.Equal(t, `{"x":2}`, wire.tools[1].Function.Arguments) - } else { - assertProtocolChatFailed(t, c, wire, err) - } - }) - } -} - func TestResponsesAgentMessageRoutingPreservesChatEligibility(t *testing.T) { for _, tc := range []struct { name, body string diff --git a/backend/internal/service/scheduler_cache.go b/backend/internal/service/scheduler_cache.go index d507d6ce2ad2..5e5dbed1643b 100644 --- a/backend/internal/service/scheduler_cache.go +++ b/backend/internal/service/scheduler_cache.go @@ -76,6 +76,8 @@ func ParseSchedulerBucket(raw string) (SchedulerBucket, bool) { // SchedulerCache 负责调度快照与账号快照的缓存读写。 type SchedulerCache interface { // GetSnapshot 读取快照并返回命中与否（ready + active + 数据完整）。 + // 返回的账号对象由本次调用独占，可以修改其顶层字段；map、切片底层数组与指针指向的值 + // 与其它调用共享，不得原地修改。 GetSnapshot(ctx context.Context, bucket SchedulerBucket) ([]*Account, bool, error) // CaptureBucketWriteToken captures the current open epoch without changing // retirement state. A tombstoned bucket returns ErrSchedulerBucketRetired. diff --git a/backend/internal/service/setting_gateway_runtime.go b/backend/internal/service/setting_gateway_runtime.go index fe5904f28635..26637f0a8211 100644 --- a/backend/internal/service/setting_gateway_runtime.go +++ b/backend/internal/service/setting_gateway_runtime.go @@ -98,7 +98,7 @@ const antigravityUserAgentVersionErrorTTL = 5 * time.Second const antigravityUserAgentVersionDBTimeout = 5 * time.Second // DefaultOpenAICodexUserAgent 是 OpenAI Codex 默认 User-Agent，用于规避浏览器 UA 的质询。 -// 默认采用 codex-tui 身份，版本段随 codexCLIVersion 一起更新。 +// 默认采用 codex_cli_rs 身份，避免历史 codex-tui 出口触发 sub2api 的身份风控指纹。 const DefaultOpenAICodexUserAgent = codexCLIUserAgent // cachedOpenAICodexUserAgent 缓存 OpenAI Codex UA（进程内缓存，60s TTL） @@ -459,7 +459,7 @@ func (s *SettingService) GetOpenAICodexUserAgent(ctx context.Context) string { // GetOpenAICodexClientVersion 返回出站声明的 Codex 客户端版本号。 // 优先级：管理员在面板覆写的版本 → 自动同步到的官方最新稳定版 → 内置常量。 -// 上游在容量紧张时按客户端身份分优先级降载，陈旧版本会被优先丢弃，故该值需保持跟随官方发布； +// 上游会按客户端身份进行风控分层，陈旧版本可能被优先丢弃，故该值需保持跟随官方发布； // 自动同步让运维不必为了跟版本而发新版本。 func (s *SettingService) GetOpenAICodexClientVersion(ctx context.Context) string { fallback := codexCLIVersion @@ -614,11 +614,11 @@ func (s *SettingService) InvalidateClaudeCodeClientVersionCache() { } // GetOpenAICodexCanonicalUserAgent 返回出站规范 Codex User-Agent。 -// 未填面板 UA 时按当前生效的客户端版本号拼出标准 Codex TUI UA。 +// 未填面板 UA 时按当前生效的客户端版本号拼出标准 Codex CLI UA。 // // 面板 UA 只贡献客户端名与 OS / 架构 / 终端指纹，版本段一律用生效版本重建：该输入框是 // 唯一能改 UA 后缀的地方，但它填写于某个历史版本，逐字沿用会把出站身份永久钉死在陈旧 -// 版本上并绕过自动同步——而陈旧身份正是上游优先降载的那一侧。 +// 版本上并绕过自动同步——而陈旧身份正是上游风控更严格的那一侧。 // 需要固定版本请填「Codex 客户端版本号」并关闭自动同步。 func (s *SettingService) GetOpenAICodexCanonicalUserAgent(ctx context.Context) string { if s == nil { diff --git a/backend/internal/service/settings_view.go b/backend/internal/service/settings_view.go index 8e27fc82399a..5ba40e8bf50f 100644 --- a/backend/internal/service/settings_view.go +++ b/backend/internal/service/settings_view.go @@ -252,7 +252,7 @@ type SystemSettings struct { EnableClientDatelineNormalization bool // 是否对 Anthropic OAuth/SetupToken 请求体做客户端 dateline 归一化（默认 true） RewriteMessageCacheControl bool // 是否改写 messages[*].content[*].cache_control（默认 false） AntigravityUserAgentVersion string // Antigravity 上游 User-Agent 版本号；空值使用配置/默认值 - OpenAICodexUserAgent string // OpenAI Codex 上游完整 User-Agent；空值由 Codex 客户端版本号拼出标准 TUI UA + OpenAICodexUserAgent string // OpenAI Codex 上游完整 User-Agent；空值由 Codex 客户端版本号拼出标准 CLI UA OpenAICodexClientVersion string // 出站声明的 Codex 客户端版本号（管理员覆写）；空值跟随自动同步值 OpenAICodexClientVersionSynced string // 自动同步到的官方最新稳定版版本号（只读展示） OpenAICodexVersionAutoSyncEnabled bool // 是否启用 Codex 客户端版本号自动同步（默认 true） diff --git a/backend/internal/service/single_account_retry_probe_test.go b/backend/internal/service/single_account_retry_probe_test.go new file mode 100644 index 000000000000..660762b908b8 --- /dev/null +++ b/backend/internal/service/single_account_retry_probe_test.go @@ -0,0 +1,191 @@ +//go:build unit + +package service + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/stretchr/testify/require" +) + +func TestIsSingleAccountRetry_ProbeEvaluatedLazilyOnce(t *testing.T) { + calls := 0 + ctx := WithSingleAccountRetryProbe(context.Background(), func(context.Context) bool { + calls++ + return true + }) + require.Zero(t, calls, "挂载时不应判定") + + require.True(t, isSingleAccountRetry(ctx)) + require.True(t, isSingleAccountRetry(ctx)) + require.Equal(t, 1, calls, "同一请求内只判定一次") +} + +func TestIsSingleAccountRetry_ProbeFalse(t *testing.T) { + ctx := WithSingleAccountRetryProbe(context.Background(), func(context.Context) bool { return false }) + require.False(t, isSingleAccountRetry(ctx)) +} + +func TestIsSingleAccountRetry_ExplicitFlagOverridesProbe(t *testing.T) { + calls := 0 + probe := func(context.Context) bool { + calls++ + return false + } + + ctx := WithSingleAccountRetryProbe(context.Background(), probe) + ctx = WithSingleAccountRetry(ctx, true, false) + require.True(t, isSingleAccountRetry(ctx)) + + legacy := context.WithValue(WithSingleAccountRetryProbe(context.Background(), probe), ctxkey.SingleAccountRetry, true) + require.True(t, isSingleAccountRetry(legacy)) + + explicitFalse := WithSingleAccountRetry(WithSingleAccountRetryProbe(context.Background(), func(context.Context) bool { + calls++ + return true + }), false, false) + require.False(t, isSingleAccountRetry(explicitFalse)) + + require.Zero(t, calls, "显式标记存在时不应触发判定") +} + +func TestIsSingleAccountRetry_ProbeIgnoresCanceledContext(t *testing.T) { + ctx, cancel := context.WithCancel(WithSingleAccountRetryProbe(context.Background(), func(probeCtx context.Context) bool { + return probeCtx.Err() == nil + })) + cancel() + require.True(t, isSingleAccountRetry(ctx)) +} + +func TestIsSingleAccountRetry_ProbeSurvivesMetadataUpdates(t *testing.T) { + ctx := WithSingleAccountRetryProbe(context.Background(), func(context.Context) bool { return true }) + ctx = WithAccountSwitchCount(ctx, 1, false) + ctx = WithThinkingEnabled(ctx, true, false) + require.True(t, isSingleAccountRetry(ctx)) +} + +// groupRecordingAccountRepo 记录单平台分组查询的分组与平台。 +type groupRecordingAccountRepo struct { + *mockAccountRepoForPlatform + groupIDs []int64 + platforms []string +} + +func (r *groupRecordingAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) { + r.groupIDs = append(r.groupIDs, groupID) + r.platforms = append(r.platforms, platform) + return r.mockAccountRepoForPlatform.ListSchedulableByGroupIDAndPlatform(ctx, groupID, platform) +} + +func newSingleAntigravityProbeTestService(accounts ...Account) (*GatewayService, *groupRecordingAccountRepo) { + repo := &groupRecordingAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{accounts: accounts}} + return &GatewayService{accountRepo: repo}, repo +} + +func probeTestAccount(id int64, platform string) Account { + return Account{ID: id, Platform: platform, Status: StatusActive, Schedulable: true} +} + +func TestWithSingleAntigravityAccountGroupProbe_QueriesOnlyWhenRead(t *testing.T) { + svc, repo := newSingleAntigravityProbeTestService(probeTestAccount(1, PlatformAntigravity)) + groupID := int64(42) + + ctx := svc.WithSingleAntigravityAccountGroupProbe(context.Background(), &groupID) + require.Empty(t, repo.groupIDs, "挂载时不应查询账号") + + require.True(t, isSingleAccountRetry(ctx)) + require.True(t, isSingleAccountRetry(ctx)) + require.Equal(t, []int64{42}, repo.groupIDs) + require.Equal(t, []string{PlatformAntigravity}, repo.platforms) +} + +func TestWithSingleAntigravityAccountGroupProbe_UsesGroupAtMountTime(t *testing.T) { + svc, repo := newSingleAntigravityProbeTestService(probeTestAccount(1, PlatformAntigravity)) + groupID := int64(42) + + ctx := svc.WithSingleAntigravityAccountGroupProbe(context.Background(), &groupID) + groupID = 99 + + require.True(t, isSingleAccountRetry(ctx)) + require.Equal(t, []int64{42}, repo.groupIDs) +} + +func TestWithSingleAntigravityAccountGroupProbe_CountsOnlyAntigravityAccounts(t *testing.T) { + groupID := int64(42) + tests := []struct { + name string + accounts []Account + want bool + }{ + {name: "no antigravity accounts", accounts: []Account{probeTestAccount(1, PlatformGemini)}, want: false}, + {name: "single antigravity account", accounts: []Account{probeTestAccount(1, PlatformAntigravity), probeTestAccount(2, PlatformGemini)}, want: true}, + {name: "two antigravity accounts", accounts: []Account{probeTestAccount(1, PlatformAntigravity), probeTestAccount(2, PlatformAntigravity)}, want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + svc, _ := newSingleAntigravityProbeTestService(tt.accounts...) + ctx := svc.WithSingleAntigravityAccountGroupProbe(context.Background(), &groupID) + require.Equal(t, tt.want, isSingleAccountRetry(ctx)) + }) + } +} + +func runRateLimitedPreCheckWithProbe(t *testing.T, single bool) (*antigravityRetryLoopResult, *recordingOKUpstream, error) { + t.Helper() + upstream := &recordingOKUpstream{} + account := &Account{ + ID: 30, + Type: AccountTypeOAuth, + Platform: PlatformAntigravity, + Schedulable: true, + Status: StatusActive, + Concurrency: 1, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + "claude-sonnet-4-5": map[string]any{ + "rate_limit_reset_at": time.Now().Add(30 * time.Second).Format(time.RFC3339), + }, + }, + }, + } + ctx := WithSingleAccountRetryProbe(context.Background(), func(context.Context) bool { return single }) + svc := &AntigravityGatewayService{} + result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: "[test]", + account: account, + accessToken: "token", + action: "generateContent", + body: []byte(`{"input":"test"}`), + httpUpstream: upstream, + requestedModel: "claude-sonnet-4-5", + handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { + return nil + }, + }) + return result, upstream, err +} + +func TestAntigravityRetryLoop_PreCheck_ProbeSingleAccount_SkipsRateLimit(t *testing.T) { + result, upstream, err := runRateLimitedPreCheckWithProbe(t, true) + + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.resp) + require.Equal(t, http.StatusOK, result.resp.StatusCode) + require.Equal(t, 1, upstream.calls) +} + +func TestAntigravityRetryLoop_PreCheck_ProbeMultiAccount_SwitchesOnRateLimit(t *testing.T) { + result, upstream, err := runRateLimitedPreCheckWithProbe(t, false) + + require.Nil(t, result) + var switchErr *AntigravityAccountSwitchError + require.ErrorAs(t, err, &switchErr) + require.Equal(t, "claude-sonnet-4-5", switchErr.RateLimitedModel) + require.Zero(t, upstream.calls) +} diff --git a/backend/internal/service/usage_billing.go b/backend/internal/service/usage_billing.go index e058ecb1b029..447260ecb451 100644 --- a/backend/internal/service/usage_billing.go +++ b/backend/internal/service/usage_billing.go @@ -42,6 +42,11 @@ type UsageBillingCommand struct { APIKeyQuotaCost float64 APIKeyRateLimitCost float64 AccountQuotaCost float64 + + // AccountQuotaHasDailyLimit / AccountQuotaHasWeeklyLimit 记录构造命令时账号是否设置了日/周上限， + // 只用于选择账号配额累加语句，不参与幂等指纹；实际累加哪些维度以数据库中的账号配置为准。 + AccountQuotaHasDailyLimit bool + AccountQuotaHasWeeklyLimit bool } func (c *UsageBillingCommand) Normalize() { diff --git a/backend/migrations/246_enable_gemini_image_generation_groups.sql b/backend/migrations/246_enable_gemini_image_generation_groups.sql new file mode 100644 index 000000000000..ae37f206fce3 --- /dev/null +++ b/backend/migrations/246_enable_gemini_image_generation_groups.sql @@ -0,0 +1,6 @@ +-- gemini / antigravity 分组的生图开关同时约束 Gemini 系入口，存量分组统一置为开启。 +-- composite 分组的同一开关还约束 OpenAI 生图，不在本迁移范围内。 +UPDATE groups +SET allow_image_generation = true +WHERE platform IN ('gemini', 'antigravity') + AND allow_image_generation = false; diff --git a/backend/migrations/gemini_image_generation_groups_migration_test.go b/backend/migrations/gemini_image_generation_groups_migration_test.go new file mode 100644 index 000000000000..81630a151bcb --- /dev/null +++ b/backend/migrations/gemini_image_generation_groups_migration_test.go @@ -0,0 +1,24 @@ +package migrations + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestEnableGeminiImageGenerationGroupsMigration(t *testing.T) { + content, err := FS.ReadFile("246_enable_gemini_image_generation_groups.sql") + require.NoError(t, err) + + var statements []string + for _, line := range strings.Split(string(content), "\n") { + if !strings.HasPrefix(strings.TrimSpace(line), "--") { + statements = append(statements, line) + } + } + sql := strings.Join(strings.Fields(strings.Join(statements, " ")), " ") + require.Equal(t, + "UPDATE groups SET allow_image_generation = true WHERE platform IN ('gemini', 'antigravity') AND allow_image_generation = false;", + sql) +} diff --git a/backend/resources/model-pricing/model_prices_and_context_window.json b/backend/resources/model-pricing/model_prices_and_context_window.json index 1e974233bafe..c6ca773cf2b3 100644 --- a/backend/resources/model-pricing/model_prices_and_context_window.json +++ b/backend/resources/model-pricing/model_prices_and_context_window.json @@ -64,6 +64,18 @@ "output_cost_per_token_above_200k_tokens": 1.2e-05, "supports_prompt_caching": true }, + "grok-4.7-fast": { + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_200k_tokens": 1.5e-06, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "litellm_provider": "xai", + "max_input_tokens": 500000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "supports_prompt_caching": true + }, "gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_batches": 1.25e-06, diff --git a/frontend/package.json b/frontend/package.json index 7b891ca2972d..c8a9d03d9975 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -23,7 +23,7 @@ "@vueuse/core": "^10.7.0", "axios": "^1.20.0", "chart.js": "^4.4.1", - "dompurify": "^3.4.14", + "dompurify": "^3.4.17", "driver.js": "^1.4.0", "file-saver": "^2.0.5", "marked": "^17.0.1", @@ -65,10 +65,9 @@ "js-cookie": "3.0.7", "form-data@<4.0.6": ">=4.0.6", "postcss@<8.5.18": ">=8.5.18", - "nanoid@<3.3.18": "3.3.18", + "dompurify@<3.4.17": ">=3.4.17", "source-map-js@<1.2.2": ">=1.2.2", - "vue@<3.5.42": ">=3.5.42", - "dompurify@<3.4.14": ">=3.4.14" + "follow-redirects@<1.16.1": ">=1.16.1" } } } diff --git a/frontend/pnpm-lock.yaml b/frontend/pnpm-lock.yaml index b6326ac8d345..d920a9e61571 100644 --- a/frontend/pnpm-lock.yaml +++ b/frontend/pnpm-lock.yaml @@ -8,10 +8,9 @@ overrides: js-cookie: 3.0.7 form-data@<4.0.6: '>=4.0.6' postcss@<8.5.18: '>=8.5.18' - nanoid@<3.3.18: 3.3.18 + dompurify@<3.4.17: '>=3.4.17' source-map-js@<1.2.2: '>=1.2.2' - vue@<3.5.42: '>=3.5.42' - dompurify@<3.4.14: '>=3.4.14' + follow-redirects@<1.16.1: '>=1.16.1' importers: @@ -39,8 +38,8 @@ importers: specifier: ^4.4.1 version: 4.5.1 dompurify: - specifier: ^3.4.14 - version: 3.4.14 + specifier: ^3.4.17 + version: 3.4.17 driver.js: specifier: ^1.4.0 version: 1.4.0 @@ -1423,7 +1422,7 @@ packages: '@tanstack/vue-virtual@3.13.23': resolution: {integrity: sha512-b5jPluAR6U3eOq6GWAYSpj3ugnAIZgGR0e6aGAgyRse0Yu6MVQQ0ZWm9SArSXWtageogn6bkVD8D//c4IjW3xQ==} peerDependencies: - vue: '>=3.5.42' + vue: ^2.7.0 || ^3.0.0 '@types/d3-array@3.2.2': resolution: {integrity: sha512-hOLWVbm7uRza0BYXpIIW5pxfrKe0W+D5lrFiAEYR+pb6w3N2SwSMaJbXdUfSEv+dT4MfHBLtn5js0LAWaO6otw==} @@ -1660,7 +1659,7 @@ packages: engines: {node: ^18.0.0 || >=20.0.0} peerDependencies: vite: ^5.0.0 || ^6.0.0 - vue: '>=3.5.42' + vue: ^3.2.25 '@vitest/coverage-v8@2.1.9': resolution: {integrity: sha512-Z2cOr0ksM00MpEfyVE8KXIYPEcBFxdbLSs56L8PO0QQMxt/6bDj45uQfxoc96v05KW3clk7vvgP0qfDit9DmfQ==} @@ -2311,8 +2310,8 @@ packages: resolution: {integrity: sha512-yS+Q5i3hBf7GBkd4KG8a7eBNNWNGLTaEwwYWUijIYM7zrlYDM0BFXHjjPWlWZ1Rg7UaddZeIDmi9jF3HmqiQ2w==} engines: {node: '>=6.0.0'} - dompurify@3.4.14: - resolution: {integrity: sha512-dVoH9z+MY+C9IilgGCk3YfFqjLi3fChm2OiKJMzh6axrJ5qwxqWaZamgmHrpv22CN/KdbZJuGEGgfQoL00LTdg==} + dompurify@3.4.17: + resolution: {integrity: sha512-B/pSlYVHQ58kqISD/yDjuJ1yfzI34m1gF4lAod2bwSRhiN4K5Cl3H8Af7hqB2PmTL911tdEudnnXn5FnhRoz0w==} driver.js@1.4.0: resolution: {integrity: sha512-Gm64jm6PmcU+si21sQhBrTAM1JvUrR0QhNmjkprNLxohOBzul9+pNHXgQaT9lW84gwg9GMLB3NZGuGolsz5uew==} @@ -2545,8 +2544,8 @@ packages: flatted@3.3.3: resolution: {integrity: sha512-GX+ysw4PBCz0PzosHDepZGANEuFCMLrnRTiEy9McGjmkCQYwRq4A/X786G/fjM/+OjsWSU1ZrY5qyARZmO/uwg==} - follow-redirects@1.16.0: - resolution: {integrity: sha512-y5rN/uOsadFT/JfYwhxRS5R7Qce+g3zG97+JrtFZlC9klX/W5hD7iiLzScI4nZqUS7DNUdhPgw4xI8W2LuXlUw==} + follow-redirects@1.16.1: + resolution: {integrity: sha512-FNvFGzoMLWmE6Yj9spb/zjd7yiNCHiAW9/Tg9CXrQ8wuu32HtlJOwWO11OJafl5FfY3DxTdQ0vj42zU1kvv5jg==} engines: {node: '>=4.0'} peerDependencies: debug: '*' @@ -3474,7 +3473,7 @@ packages: resolution: {integrity: sha512-khUlZSwt9xXCaTbbxFYBKDc/bWAGWJjOgvxETwkTN7KRm66EeT1ZdZj6i2ceh9sP2Pzqsbc704r2yngBrxBVug==} peerDependencies: typescript: '>=4.4.4' - vue: '>=3.5.42' + vue: ^2.7.0 || ^3.5.11 peerDependenciesMeta: typescript: optional: true @@ -3924,7 +3923,7 @@ packages: peerDependencies: react: ^19.0.0 solid-js: ^1.9.0 - vue: '>=3.5.42' + vue: ^3.2.0 peerDependenciesMeta: react: optional: true @@ -4356,7 +4355,7 @@ packages: resolution: {integrity: sha512-jqxtL8KZ6YJ5NTv6XzrzLS7osyegOi28UGNZW0h9OkDL7Sh1396ht4Dorh04aKrl2LiSalQ84WtqiG0RIJb0tA==} peerDependencies: chart.js: ^4.1.1 - vue: '>=3.5.42' + vue: ^3.0.0-0 || ^2.7.0 vue-component-type-helpers@2.2.12: resolution: {integrity: sha512-YbGqHZ5/eW4SnkPNR44mKVc6ZKQoRs/Rux1sxC6rdwXb4qpbOSYfDr9DsTHolOTGmIKgM9j141mZbBeg05R1pw==} @@ -4367,7 +4366,7 @@ packages: hasBin: true peerDependencies: '@vue/composition-api': ^1.0.0-rc.1 - vue: '>=3.5.42' + vue: ^3.0.0-0 || ^2.6.0 peerDependenciesMeta: '@vue/composition-api': optional: true @@ -4392,12 +4391,12 @@ packages: engines: {node: '>= 16'} deprecated: v9 and v10 no longer supported. please migrate to v11. about maintenance status, see https://vue-i18n.intlify.dev/guide/maintenance.html peerDependencies: - vue: '>=3.5.42' + vue: ^3.0.0 vue-router@4.6.4: resolution: {integrity: sha512-Hz9q5sa33Yhduglwz6g9skT8OBPii+4bFn88w6J+J4MfEo4KRRpmiNG/hHHkdbRFlLBOqxN8y8gf2Fb0MTUgVg==} peerDependencies: - vue: '>=3.5.42' + vue: ^3.5.0 vue-tsc@2.2.12: resolution: {integrity: sha512-P7OP77b2h/Pmk+lZdJ0YWs+5tJ6J2+uOQPo7tlBnY44QqQSPYvS0qVT4wqDJgwrZaLe47etJLLQRFia71GYITw==} @@ -6021,7 +6020,7 @@ snapshots: '@types/dompurify@3.2.0': dependencies: - dompurify: 3.4.14 + dompurify: 3.4.17 '@types/estree-jsx@1.0.5': dependencies: @@ -6520,7 +6519,7 @@ snapshots: axios@1.20.0: dependencies: - follow-redirects: 1.16.0 + follow-redirects: 1.16.1 form-data: 4.0.6 https-proxy-agent: 5.0.1 proxy-from-env: 2.1.0 @@ -6959,7 +6958,7 @@ snapshots: dependencies: esutils: 2.0.3 - dompurify@3.4.14: + dompurify@3.4.17: optionalDependencies: '@types/trusted-types': 2.0.7 @@ -7253,7 +7252,7 @@ snapshots: flatted@3.3.3: {} - follow-redirects@1.16.0: {} + follow-redirects@1.16.1: {} for-in@1.0.2: {} @@ -8058,7 +8057,7 @@ snapshots: d3-sankey: 0.12.3 dagre-d3-es: 7.0.14 dayjs: 1.11.20 - dompurify: 3.4.14 + dompurify: 3.4.17 katex: 0.16.45 khroma: 2.1.0 lodash-es: 4.18.1 diff --git a/frontend/src/components/account/AccountCapacityCell.vue b/frontend/src/components/account/AccountCapacityCell.vue index a1873f36011d..271bfeb1b1f8 100644 --- a/frontend/src/components/account/AccountCapacityCell.vue +++ b/frontend/src/components/account/AccountCapacityCell.vue @@ -39,6 +39,7 @@ import { computed } from 'vue' import { useI18n } from 'vue-i18n' import type { Account } from '@/types' +import { supportsAccountQuota } from '@/utils/accountQuota' import CapacityBadge from '@/components/account/CapacityBadge.vue' import QuotaBadge from '@/components/account/QuotaBadge.vue' @@ -177,7 +178,7 @@ const formatCost = (value: number | null | undefined) => { } // ====== 配额 ====== -const isQuotaEligible = computed(() => props.account.type === 'apikey' || props.account.type === 'bedrock') +const isQuotaEligible = computed(() => supportsAccountQuota(props.account.type)) const showDailyQuota = computed(() => isQuotaEligible.value && props.account.quota_daily_limit != null && props.account.quota_daily_limit > 0 diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index 22c840209f2a..6d9d3f668082 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -585,7 +585,7 @@ />
{{ t('admin.accounts.gemini.quotaPolicy.title') }}
{{ t('admin.accounts.gemini.quotaPolicy.note') }}
@@ -608,7 +608,7 @@ v-if="showGeminiTodayStats && todayStats" class="mb-0.5 flex items-center" > -
+
{{ formatKeyRequests }} req @@ -635,14 +635,37 @@
-
+ + +
-
+
{{ error }}
@@ -650,7 +673,6 @@ {{ t('admin.accounts.gemini.rateLimit.unlimited') }}
- - -
@@ -823,6 +822,7 @@ import { buildOpenAIUsageRefreshKey } from '@/utils/accountUsageRefresh' import { enqueueUsageRequest } from '@/utils/usageLoadQueue' import { formatCompactNumber, formatRelativeTime } from '@/utils/format' import { isOfficialGeminiBaseUrl } from '@/utils/geminiExplicitCache' +import { supportsAccountQuota } from '@/utils/accountQuota' import UsageProgressBar from './UsageProgressBar.vue' import AccountQuotaInfo from './AccountQuotaInfo.vue' import ClaudeResetCreditsCell from './ClaudeResetCreditsCell.vue' @@ -1875,7 +1875,7 @@ const makeQuotaBar = ( } const hasApiKeyQuota = computed(() => { - if (props.account.type !== 'apikey' && props.account.type !== 'bedrock') return false + if (!supportsAccountQuota(props.account.type)) return false return ( (props.account.quota_daily_limit ?? 0) > 0 || (props.account.quota_weekly_limit ?? 0) > 0 || diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index c6967eb82575..df8c394d8e92 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -2412,9 +2412,9 @@ />
- +
@@ -4279,6 +4279,7 @@ import { } from '@/utils/format' import { createStableObjectKeyResolver } from '@/utils/stableObjectKey' import { supportsUpstreamModelSync } from '@/utils/upstreamModelSync' +import { supportsAccountQuota } from '@/utils/accountQuota' import { getAccountExpiryTimestamp } from '@/components/account/accountExpiry' import { VERTEX_LOCATION_OPTIONS } from '@/constants/account' import { @@ -6447,9 +6448,9 @@ const createAccountAndFinish = async ( if (!applyTempUnschedConfig(credentials)) { return } - // Inject quota limits for apikey/bedrock accounts + // Inject quota limits for apikey/bedrock/service_account accounts let finalExtra = withUpstreamRequestIdHeader(extra) - if (type === 'apikey' || type === 'bedrock') { + if (supportsAccountQuota(type)) { const quotaExtra: Record = { ...(finalExtra || {}) } if (editQuotaLimit.value != null && editQuotaLimit.value > 0) { quotaExtra.quota_limit = editQuotaLimit.value diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index 3580d39df4bf..90c381e325c9 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -2401,9 +2401,9 @@ @update:quotaNotifyTotalThresholdType="quotaNotifyState.total.thresholdType = $event" />
- +
@@ -3405,6 +3405,7 @@ import { extractApiErrorMessage, extractI18nErrorMessage } from '@/utils/apiErro import { createStableObjectKeyResolver } from '@/utils/stableObjectKey' import { supportsUpstreamModelSync } from '@/utils/upstreamModelSync' import { takesSystemPrompt } from '@/utils/systemPromptBinding' +import { supportsAccountQuota } from '@/utils/accountQuota' import { getAccountExpiryTimestamp } from '@/components/account/accountExpiry' import { allSelectedGroupsEnableLongContextPricing } from '@/components/account/longContextBilling' import { VERTEX_LOCATION_OPTIONS } from '@/constants/account' @@ -4513,8 +4514,8 @@ const syncFormFromAccount = (newAccount: Account | null) => { } } - // Load quota limit for apikey/bedrock accounts (bedrock quota is also loaded in its own branch above) - if (newAccount.type === 'apikey' || newAccount.type === 'bedrock') { + // Load quota limit for apikey/bedrock/service_account accounts (bedrock quota is also loaded in its own branch above) + if (supportsAccountQuota(newAccount.type)) { const quotaVal = extra?.quota_limit as number | undefined editQuotaLimit.value = (quotaVal && quotaVal > 0) ? quotaVal : null const dailyVal = extra?.quota_daily_limit as number | undefined @@ -6136,8 +6137,8 @@ const handleSubmit = async () => { updatePayload.extra = newExtra } - // For apikey/bedrock accounts, handle quota_limit in extra - if (props.account.type === 'apikey' || props.account.type === 'bedrock') { + // For apikey/bedrock/service_account accounts, handle quota_limit in extra + if (supportsAccountQuota(props.account.type)) { const currentExtra = (updatePayload.extra as Record) || (props.account.extra as Record) || {} const newExtra: Record = { ...currentExtra } diff --git a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts index bd188e522870..05b67e41931b 100644 --- a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts +++ b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts @@ -1780,6 +1780,122 @@ describe('AccountUsageCell', () => { expect(getUsage).not.toHaveBeenCalled() }) + it('Vertex 账号不设上限时不显示配额进度条', async () => { + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ + id: 4011, + platform: 'gemini', + type: 'service_account', + credentials: { tier_id: 'vertex', project_id: 'vertex-proj', location: 'global' }, + extra: {} + }), + todayStats: { + requests: 3, + tokens: 1200, + cost: 0.42, + standard_cost: 0.42, + user_cost: 0.5 + } + }, + global: { + stubs: { + UsageProgressBar: true, + AccountQuotaInfo: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('A $0.42') + expect(wrapper.findAll('usage-progress-bar-stub')).toHaveLength(0) + expect(getUsage).not.toHaveBeenCalled() + }) + + it('Vertex 账号设置了配额上限时显示配额进度条', async () => { + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ + id: 4012, + platform: 'gemini', + type: 'service_account', + credentials: { tier_id: 'vertex', project_id: 'vertex-proj', location: 'global' }, + extra: { quota_daily_start: new Date().toISOString() }, + quota_used: 123.456, + quota_limit: 500, + quota_daily_limit: 50, + quota_daily_used: 12.5, + quota_weekly_limit: 100, + quota_weekly_used: 25 + }), + todayStats: null, + todayStatsLoading: false + }, + global: { + stubs: { + UsageProgressBar: true, + AccountQuotaInfo: true + } + } + }) + + await flushPromises() + + const bars = wrapper.findAll('usage-progress-bar-stub') + expect(bars.map((bar) => bar.attributes('label'))).toEqual(['1d', '7d', 'total']) + expect(getUsage).not.toHaveBeenCalled() + }) + + it('Gemini 刷新时显示加载或错误状态，不以旧用量覆盖官方状态展示', async () => { + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ + id: 4013, + platform: 'gemini', + type: 'apikey', + credentials: { + base_url: 'https://generativelanguage.googleapis.com', + api_key: 'AIza-official', + tier_id: 'aistudio_paid' + } + }), + requestBatchedUsage: vi.fn(), + batchedUsage: { + updated_at: null, + five_hour: null, + seven_day: null, + seven_day_sonnet: null, + gemini_pro_daily: { + utilization: 20, + resets_at: '2026-10-11T07:00:00Z', + remaining_seconds: 3600 + } + } + }, + global: { + stubs: { + UsageProgressBar: true, + AccountQuotaInfo: true + } + } + }) + await flushPromises() + expect(wrapper.findAll('usage-progress-bar-stub')).toHaveLength(1) + + await wrapper.setProps({ batchedUsageLoading: true }) + expect(wrapper.findAll('usage-progress-bar-stub')).toHaveLength(0) + expect(wrapper.find('.animate-pulse').exists()).toBe(true) + + await wrapper.setProps({ batchedUsageLoading: false, batchedUsageError: 'quota refresh failed' }) + expect(wrapper.findAll('usage-progress-bar-stub')).toHaveLength(0) + expect(wrapper.text()).toContain('quota refresh failed') + + await wrapper.setProps({ batchedUsageError: null }) + expect(wrapper.findAll('usage-progress-bar-stub')).toHaveLength(1) + wrapper.unmount() + }) + it('Gemini 第三方上游 API Key 按普通 Key 账号展示今日统计与额度，不显示 Gemini 徽章', async () => { const wrapper = mount(AccountUsageCell, { props: { diff --git a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts index 2708133b96e1..0c6773c635ad 100644 --- a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts @@ -2254,6 +2254,26 @@ describe('EditAccountModal', () => { expect(updateAccountMock).toHaveBeenCalledTimes(1) }) + it('shows quota limit card for Vertex SA account and keeps quota limits on save', async () => { + const account = buildVertexAccount() + account.extra = { quota_limit: 200, quota_daily_limit: 20, quota_used: 35.5 } + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + + const wrapper = mountModal(account) + expect(wrapper.findComponent({ name: 'QuotaLimitCard' }).exists()).toBe(true) + + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + const extra = updateAccountMock.mock.calls[0]?.[1]?.extra + expect(extra?.quota_limit).toBe(200) + expect(extra?.quota_daily_limit).toBe(20) + expect(extra?.quota_weekly_limit).toBeUndefined() + }) + it('blocks Vertex SA save when neither credentials_status nor legacy json indicates existence', async () => { const account = buildVertexAccount() account.credentials = { diff --git a/frontend/src/components/admin/account/AccountActionMenu.vue b/frontend/src/components/admin/account/AccountActionMenu.vue index a7eb904db22b..168d45077e69 100644 --- a/frontend/src/components/admin/account/AccountActionMenu.vue +++ b/frontend/src/components/admin/account/AccountActionMenu.vue @@ -78,6 +78,7 @@ import { Icon } from '@/components/icons' import type { Account } from '@/types' import BaseDialog from '@/components/common/BaseDialog.vue' import { takesSystemPrompt } from '@/utils/systemPromptBinding' +import { supportsAccountQuota } from '@/utils/accountQuota' const props = defineProps<{ show: boolean; account: Account | null; anchorRect: DOMRect | null; busy?: boolean }>() const emit = defineEmits(['close', 'test', 'stats', 'schedule', 'duplicate', 'reauth', 'refresh-token', 'recover-state', 'reset-quota', 'set-privacy', 'create-spark-shadow']) @@ -152,7 +153,7 @@ const isShadow = computed(() => props.account?.parent_account_id != null) const isOpenAIOAuthParent = computed(() => isOpenAIOAuth.value && !isShadow.value) const supportsPrivacy = computed(() => (isAntigravityOAuth.value || isOpenAIOAuth.value) && !isShadow.value) const hasQuotaLimit = computed(() => { - return (props.account?.type === 'apikey' || props.account?.type === 'bedrock') && ( + return supportsAccountQuota(props.account?.type) && ( (props.account?.quota_limit ?? 0) > 0 || (props.account?.quota_daily_limit ?? 0) > 0 || (props.account?.quota_weekly_limit ?? 0) > 0 diff --git a/frontend/src/components/admin/account/AccountConsoleFilters.vue b/frontend/src/components/admin/account/AccountConsoleFilters.vue index 21feb08a9daa..b6050298e53e 100644 --- a/frontend/src/components/admin/account/AccountConsoleFilters.vue +++ b/frontend/src/components/admin/account/AccountConsoleFilters.vue @@ -5,6 +5,22 @@
+